evaluation.py

import pytorch_lightning as pl
import torch

torch.set_float32_matmul_precision('medium')
from pytorch_lightning.loggers import WandbLogger
from torch.utils.data import DataLoader
from models import build_model
from datasets import build_dataset
from utils.utils import set_seed
import hydra
from omegaconf import OmegaConf
from torch.utils.data import Subset
from torchsummary import summary

import wandb

class SubsetWithAttributes(Subset):
    def __init__(self, dataset, indices):
        super().__init__(dataset, indices)
        self.__dict__.update(dataset.__dict__)
        self.collate_fn = getattr(dataset, 'collate_fn', None)

@hydra.main(version_base=None, config_path="configs", config_name="config")
def evaluation(cfg):
    set_seed(cfg.seed)
    OmegaConf.set_struct(cfg, False)  # Open the struct
    cfg = OmegaConf.merge(cfg, cfg.method)
    cfg['eval'] = True

    model = build_model(cfg)
    # summary(model, (1, 28, 28))
    print(model)
    
    # dataset modify
    # # Split the dataset into 80% training and 20% validation
    # full_dataset = build_dataset(cfg)
    # train_size = int(0.8 * len(full_dataset))
    # val_size = len(full_dataset) - train_size
    # # Create train and validation subsets
    # train_indices = list(range(train_size))
    # val_indices = list(range(train_size, len(full_dataset)))

    # train_set = SubsetWithAttributes(full_dataset, train_indices)
    # val_set = SubsetWithAttributes(full_dataset, val_indices)
    val_set = build_dataset(cfg, val=True, missing=cfg.missing)
    # val_set = build_dataset(cfg)

    eval_batch_size = max(cfg.method['eval_batch_size'] // len(cfg.devices) // val_set.data_chunk_size, 1)

    val_loader = DataLoader(
        val_set, batch_size=eval_batch_size, num_workers=cfg.load_num_workers, shuffle=False, drop_last=False,
        collate_fn=val_set.collate_fn)
    
    # disable Wandb when debug
    if cfg.debug:
        wandb_logger = None
    else:
        run = wandb.init(
            project=cfg.project_name, 
            name=cfg.exp_name
        )
        wandb_logger = WandbLogger(project="causalode", name=cfg.exp_name)

    trainer = pl.Trainer(
        inference_mode=True,
        logger=None if cfg.debug else wandb_logger,
        devices=1 if cfg.debug else cfg.devices,
        accelerator="cpu" if cfg.debug else "gpu",
        profiler="simple",
    )

    trainer.validate(model=model, dataloaders=val_loader, ckpt_path=cfg.ckpt_path)


if __name__ == '__main__':
    evaluation()