A potential bug in CheckpointLoader #2506
Description
There might a potential bug in following lines:
MONAI/monai/handlers/checkpoint_loader.py
Lines 95 to 104 in ce45640
If we
- set
CheckpointLoader.strict_shape = False - load other objects besides model's
state_dictfrom checkpoint (e.g. whenload_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)