A potential bug in CheckpointLoader

[sandylaker]

There might a potential bug in following lines:

MONAI/monai/handlers/checkpoint_loader.py

Lines 95 to 104 in ce45640

	if not self.strict_shape:
	k, _ = list(self.load_dict.items())[0]
	# single object and checkpoint is directly a state_dict
	if len(self.load_dict) == 1 and k not in checkpoint:
	checkpoint = {k: checkpoint}
	# skip items that don't match data shape
	for k, obj in self.load_dict.items():
	checkpoint[k] = copy_model_state(obj, checkpoint, inplace=False)[0]

If we

1. set CheckpointLoader.strict_shape = False
2. load other objects besides model's state_dict from checkpoint (e.g. when load_dict = {"model": model, "optimizer": optimizer}),

the obj might be model or optimizer in Line 103, while checkpoint is always a dict containing keys "model" and "optimizer". If the obj is optimizer, then this line will raise an Error since optimizer is not iterable. If the obj is model, {dst_prefix}{s}(see here) will never be matched, since s is just "model" or "optimizer". Thus, the state_dict from check point will never be loaded.

Although the conditions to trigger this error is barely fulfilled in practice (e.g. when we do transfer learning and the checkpoint happens to contain some non-iterable objects like tensors, optimizer, etc.), it is better to be aware of this issue.

You can run the following test script to see the detailed error message:

from ignite.engine import Engine
from torch.optim import SGD
import torch
import torch.nn as nn
import tempfile
from monai.handlers import CheckpointLoader
import os.path as osp


def main(output_dir):
	# build dummy data set and data loader
	x = torch.randn(10, 2)
	y = torch.randn(10)
	train_set = torch.utils.data.TensorDataset(x, y)
	train_loader = torch.utils.data.DataLoader(
		train_set, 
		batch_size=2, 
		num_workers=0, 
		shuffle=False)

	# build dummy model, optimizer and engine
	model = nn.Linear(2, 1, bias=False)
	optimizer = SGD(model.parameters(), lr=0.001)

	def update_model(engine, batch):
		x_, y_ = batch
		optimizer.zero_grad()
		loss = (model(x_) - y_).sum()
		loss.backward()
		optimizer.step()
		return loss.item()

	trainer = Engine(update_model)

	# save a checkpoint for loading
	to_save = {'model': model.state_dict(), 'optimizer': optimizer.state_dict()}
	to_load = {'model': model, 'optimizer': optimizer}
	ckpt = torch.save(to_save, osp.join(tmp_dir, "dummy.pth"))

	load_handler = CheckpointLoader(osp.join(tmp_dir, 'dummy.pth'), to_load, strict_shape=False)
	load_handler.attach(trainer)

	trainer.run(train_loader, max_epochs=1)


if __name__ == '__main__':
	with tempfile.TemporaryDirectory() as tmp_dir:
		main(tmp_dir)

[Nic-Ma]

Hi @sandylaker ,

Thanks for raising this issue, I am checking it.