Could you please provide an example about how to resume training? #64

Activity

1. GangLii commented on Sep 13, 2024

   

   GangLii

   on Sep 13, 2024

   Last edited by GangLii

   Collaborator

   More actions

   Thanks for your interest in our library. Could you share which loss function you are using? As some loss function Class also involve optimization parameters that are updated in each iteration(e.g., moving average estimator self.u_pos in AveragePrecisionLoss() ), the degraded performance might be caused by the re-initialized loss function for each resuming. So, to resume training exactly, you also need to load the previous optimization parameters in the loss function Class. Currently, you can try the naive solution below to save the previous loss function Class. We'll incorporate this feature to support resuming training in our further development.

   # e.g., loss_fn = APLoss(10, margin=1, gamma=0.1)
   if resume:
     checkpoint = torch.load(checkpoint_path)
     model.load_state_dict(checkpoint['model_state_dict'])
     optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
     start_epoch = checkpoint['epoch']
     loss_fn =  checkpoint['loss_fn']
     loss = checkpoint['loss']
     model.train()
   
   torch.save({
   'epoch': epoch + 1,
   'model_state_dict': model.module.state_dict(),
   'optimizer_state_dict': optimizer.state_dict(),
   'loss_fn': loss_fn,
   'loss': loss,
   }, checkpoint_path)
   
2. optmai

   closed this as completedon Sep 15, 2024
3. optmai

   reopened this on Mar 3, 2026