Resume training after the code crash due to changes in the dataloader

[AhmedHSalamah]

Hey,

I have embbded Dali Dataloader in this repo and I found the code crashed and I am not able to reproduce the error to solve it. It might be The only way that I over come this issue is by resuming the training and saving each epoch as shown below.

Could you please confirm that this is the right way to do and it will not affect my training process?

def save_ckp(state, is_best, checkpoint_dir, best_model_dir):
    f_path = checkpoint_dir / 'checkpoint.pt'
    torch.save(state, f_path)
    if is_best:
        best_fpath = best_model_dir / 'best_model.pt'
        shutil.copyfile(f_path, best_fpath)



def load_ckp(checkpoint_fpath, model, optimizer):
    checkpoint = torch.load(checkpoint_fpath)
    model.load_state_dict(checkpoint['state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer'])
    return model, optimizer, checkpoint['epoch']


checkpoint = {
    'epoch': epoch + 1,
    'state_dict': model.state_dict(),
    'optimizer': optimizer.state_dict()
}
# For Each epoch
save_ckp(checkpoint, is_best, checkpoint_dir, model_dir)


....
# To resume training 
model = MyModel(*args, **kwargs)
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
ckp_path = "path/to/checkpoint/checkpoint.pt"
model, optimizer, start_epoch = load_ckp(ckp_path, model, optimizer)

[yoshitomo-matsubara]

Hi @AhmedHSalamah

Could you please clarify what the actual problem is and what changes caused what errors? I don't see any torchdistill-specific lines in the code above.

On top of that, you've already asked a similar question here, using another GitHub account (correct me if I'm wrong), and the issue was resolved.