Can this model resume from former checkpoint?

[scpoppin1207]

I don't see any function or parameters to load from existed checkpoint.

[xiongxyowo]

Hi, you can add/modify the following code to the appropriate place to resume training:

# save meta-data
checkpoint = {'model': model.state_dict(), 'optim': optim.state_dict(), \
              'epoch': epoch, 'lr': scheduler.state_dict(),}
torch.save(checkpoint, path)

# load meta-data
checkpoint = torch.load(path)
model.load_state_dict(checkpoint['model'])
optim.load_state_dict(checkpoint['optim'])
epoch = checkpoint(['epoch'])
scheduler = CosineAnnealingLR(optim, args.epoch, eta_min=1.0e-7, last_epoch=epoch)

You may also refer to the pytorch tutorial.

[scpoppin1207]

Thxs 👍

[scpoppin1207]

model.load_state_dict(checkpoint['model'])
KeyError: 'model'

[scpoppin1207]

It seems that there's no key named "model"

[xiongxyowo]

If you are using the original save code, the weights of the model are not wrapped in the model:

torch.save(model.state_dict(), os.path.join(args.save_path, 'SAM2-UNet-%d.pth' % (epoch + 1)))

The new save code above wraps the state dict in the model:

checkpoint = {'model': model.state_dict(), 'optim': optim.state_dict(), \
              'epoch': epoch, 'lr': scheduler.state_dict(),}
torch.save(checkpoint, path)