In load_from_checkpoint, "TypeError: __init__ () missing 1 required positional argument:'cfg'"

[Keiku]

I am experimenting with the following repository. Keiku/PyTorch-Lightning-CIFAR10: "Not too complicated" training code for CIFAR-10 by PyTorch Lightning

It is implemented as follows. I was able to train/validation. But test has not been implemented yet.

import os

import hydra
import torch
from hydra.core.hydra_config import HydraConfig
from omegaconf import DictConfig, OmegaConf
from pytorch_lightning import Trainer, seed_everything
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger

from datamodule import LitCIFAR10DataModule
from model import LitCIFAR10Model


@hydra.main(config_path="./configs", config_name="default.yaml")
def main(cfg: DictConfig) -> None:

    if "experiments" in cfg.keys():
        cfg = OmegaConf.merge(cfg, cfg.experiments)

    seed_everything(0)
    os.environ["CUDA_VISIBLE_DEVICES"] = cfg.runs.gpu_id

    if cfg.runs.logger == "wandb":
        logger = WandbLogger(name=cfg.model.classifier, project="cifar10")
    elif cfg.runs.logger == "tensorboard":
        logger = TensorBoardLogger(cfg.train.tensorboard_dir, name=cfg.model.classifier)

    checkpoint = ModelCheckpoint(monitor="acc/val", mode="max", save_last=True)

    trainer = Trainer(
        fast_dev_run=cfg.runs.dev,
        logger=logger if not (cfg.runs.dev or cfg.runs.evaluate) else None,
        gpus=-1,
        deterministic=True,
        weights_summary=None,
        log_every_n_steps=1,
        max_epochs=cfg.train.num_epochs,
        checkpoint_callback=checkpoint,
        precision=cfg.runs.precision,
        resume_from_checkpoint=cfg.train.checkpoint,
    )

    datamodule = LitCIFAR10DataModule(cfg)
    model = LitCIFAR10Model(cfg)

    if cfg.runs.evaluate:
        model = LitCIFAR10Model.load_from_checkpoint(checkpoint_path=cfg.test.checkpoint)
        trainer.test(model, datamodule.test_dataloader())
    else:
        trainer.fit(model, datamodule)
        trainer.test()

if __name__ == "__main__":
    main()

I tried using load_from_checkpoint for test, but I get the following error. How can I solve it?

⋊> /m/n/k/c/PyTorch-Lightning-CIFAR10 on main ⨯
pipenv run python train.py +experiments=test_exp01 hydra.run.dir=outputs/test_exp01
Global seed set to 0
GPU available: True, used: True
TPU available: None, using: 0 TPU cores
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
Traceback (most recent call last):
  File "train.py", line 48, in main
    model = LitCIFAR10Model.load_from_checkpoint(checkpoint_path=cfg.test.checkpoint)
  File "/home/anasys/.local/share/virtualenvs/PyTorch-Lightning-CIFAR10-fAnnMMRx/lib/python3.8/site-packages/pytorch_lightning/core/saving.py", line 159, in load_from_checkpoint
    model = cls._load_model_state(checkpoint, strict=strict, **kwargs)
  File "/home/anasys/.local/share/virtualenvs/PyTorch-Lightning-CIFAR10-fAnnMMRx/lib/python3.8/site-packages/pytorch_lightning/core/saving.py", line 199, in _load_model_state
    model = cls(**_cls_kwargs)
TypeError: __init__() missing 1 required positional argument: 'cfg'

Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.
⋊> /m/n/k/c/PyTorch-Lightning-CIFAR10 on main ⨯

[Keiku]

I was able to infer test by setting as follows. I think the documentation is missing, so please refer to my implementation.

    if cfg.runs.evaluate:
        hparams = OmegaConf.load(cfg.test.hparams)
        model = LitCIFAR10Model.load_from_checkpoint(checkpoint_path=cfg.test.checkpoint, **hparams)
        trainer.test(model, datamodule.test_dataloader())
    else:
        trainer.fit(model, datamodule)
        trainer.test()