Error when trying to continue training a saved model

[LucaMarconato]

• inferno version:
• Python version:
• Operating System:

Description

I build a model and saved it using

trainer.save_every((1, 'epochs'))
trainer.save_to_directory(folder)

When I rerun my Python script to load and continue training the previous model I get an error.

What I Did

This is my code.

def train(load=False, folder='out'):
    print('starting training')
    os.makedirs(folder, exist_ok=True)

    # setup logger
    Logger.instance().setup('log')

    vae = Vae()

    ds = MyDataset(root_folder=root_folder, training=True)
    train_loader = torch.utils.data.DataLoader(ds, batch_size=512, num_workers=16)

    # build trainer
    trainer = Trainer(vae)
    trainer.cuda()

    trainer.build_criterion(vae.loss_function())
    trainer.build_optimizer('Adam', lr=0.001)
    # trainer.validate_every((2, 'epochs'))
    trainer.save_every((1, 'epochs'))
    trainer.save_to_directory(folder)
    trainer.set_max_num_epochs(100)

    # bind loaders
    trainer.bind_loader('train', train_loader, num_inputs=1, num_targets=1)

    # bind callbacks
    trainer.register_callback(GarbageCollection())
    # trainer.register_callback(ShowMinimalConsoleInfo())

    if load:
        trainer.load()
    trainer.fit()

When calling train(load=True) I get the following error:

  File "main.py", line 104, in my_train
    trainer.fit()
  File "/data/l989o/anaconda3/envs/hemo/lib/python3.7/site-packages/inferno/trainers/basic.py", line 1336, in fit
    self.train_for(break_callback=lambda *args: self.stop_fitting(max_num_iterations,
  File "/data/l989o/anaconda3/envs/hemo/lib/python3.7/site-packages/inferno/trainers/basic.py", line 1410, in train_for
    batch = self.fetch_next_batch('train')
  File "/data/l989o/anaconda3/envs/hemo/lib/python3.7/site-packages/inferno/trainers/basic.py", line 1092, in fetch_next_batch
    self._loader_iters.update({from_loader: self._loaders[from_loader].__iter__()})
KeyError: 'train'

Any ideas how to fix it? Thanks.

[DerThorsten]

on the first glance the code looks good, let my try to reproduce the error

[DerThorsten]

I can reproduce this:

import torch.nn as nn
from inferno.io.box.cifar import get_cifar10_loaders
from inferno.trainers.basic import Trainer
from inferno.trainers.callbacks.logging.tensorboard import TensorboardLogger
from inferno.extensions.layers.convolutional import ConvELU2D
from inferno.extensions.layers.reshape import Flatten

# Fill these in:
LOG_DIRECTORY = 'fo'
SAVE_DIRECTORY = 'fo'
DATASET_DIRECTORY = 'fo'
DOWNLOAD_CIFAR = True
USE_CUDA = True

# Build torch model
model = nn.Sequential(
    ConvELU2D(in_channels=3, out_channels=8, kernel_size=3),
    nn.MaxPool2d(kernel_size=2, stride=2),
    ConvELU2D(in_channels=8, out_channels=8, kernel_size=3),
    nn.MaxPool2d(kernel_size=2, stride=2),
    ConvELU2D(in_channels=8, out_channels=8, kernel_size=3),
    nn.MaxPool2d(kernel_size=2, stride=2),
    Flatten(),
    nn.Linear(in_features=(8 * 4 * 4), out_features=10),
    nn.LogSoftmax(dim=1)
)

# Load loaders
train_loader, validate_loader = get_cifar10_loaders(DATASET_DIRECTORY,
                                                    download=DOWNLOAD_CIFAR)

# Build trainer
trainer = Trainer(model) \
  .build_criterion('NLLLoss') \
  .build_metric('CategoricalError') \
  .build_optimizer('Adam') \
  .validate_every((2, 'epochs')) \
  .save_every((1, 'epochs')) \
  .save_to_directory(SAVE_DIRECTORY) \
  .set_max_num_epochs(10) \
  .build_logger(TensorboardLogger(log_scalars_every=(1, 'iteration'),
                                  log_images_every='never'),
                log_directory=LOG_DIRECTORY)

# Bind loaders
trainer \
    .bind_loader('train', train_loader) \
    .bind_loader('validate', validate_loader)

if USE_CUDA:
  trainer.cuda()

if False: # set this to true after a bit of inital training
  trainer.load()

# Go!
trainer.fit()

[DerThorsten]

Fix: bind the loaders after load

import torch.nn as nn
from inferno.io.box.cifar import get_cifar10_loaders
from inferno.trainers.basic import Trainer
from inferno.trainers.callbacks.logging.tensorboard import TensorboardLogger
from inferno.extensions.layers.convolutional import ConvELU2D
from inferno.extensions.layers.reshape import Flatten

# Fill these in:
LOG_DIRECTORY = 'fo'
SAVE_DIRECTORY = 'fo'
DATASET_DIRECTORY = 'fo'
DOWNLOAD_CIFAR = True
USE_CUDA = True

# Build torch model
model = nn.Sequential(
    ConvELU2D(in_channels=3, out_channels=8, kernel_size=3),
    nn.MaxPool2d(kernel_size=2, stride=2),
    ConvELU2D(in_channels=8, out_channels=8, kernel_size=3),
    nn.MaxPool2d(kernel_size=2, stride=2),
    ConvELU2D(in_channels=8, out_channels=8, kernel_size=3),
    nn.MaxPool2d(kernel_size=2, stride=2),
    Flatten(),
    nn.Linear(in_features=(8 * 4 * 4), out_features=10),
    nn.LogSoftmax(dim=1)
)

# Load loaders
train_loader, validate_loader = get_cifar10_loaders(DATASET_DIRECTORY,
                                                    download=DOWNLOAD_CIFAR)

# Build trainer
trainer = Trainer(model) \
  .build_criterion('NLLLoss') \
  .build_metric('CategoricalError') \
  .build_optimizer('Adam') \
  .validate_every((2, 'epochs')) \
  .save_every((1, 'epochs')) \
  .save_to_directory(SAVE_DIRECTORY) \
  .set_max_num_epochs(10) \
  .build_logger(TensorboardLogger(log_scalars_every=(1, 'iteration'),
                                  log_images_every='never'),
                log_directory=LOG_DIRECTORY)

if USE_CUDA:
  trainer.cuda()

if True:
  trainer.load()

# Bind loaders after load!
trainer \
    .bind_loader('train', train_loader) \
    .bind_loader('validate', validate_loader)

# Go!
trainer.fit()

[LucaMarconato]

Perfect thanks, now it works. I noticed that also callback must be assigned after calling trainer.load() otherwise they get unregistered.