I am working with CORnet-Z and I am building a separate test file.
The model seems to be saved as .pth.tar files
if FLAGS.output_path is not None: records.append(results) if len(results) > 1: pickle.dump(records, open(os.path.join(FLAGS.output_path, 'results.pkl'), 'wb')) ckpt_data = {} ckpt_data['flags'] = FLAGS.__dict__.copy() ckpt_data['epoch'] = epoch ckpt_data['state_dict'] = model.state_dict() ckpt_data['optimizer'] = trainer.optimizer.state_dict() if save_model_secs is not None: if time.time() - recent_time > save_model_secs: torch.save(ckpt_data, os.path.join(FLAGS.output_path, 'latest_checkpoint.pth.tar')) recent_time = time.time()
What would be the best approach to load this model and run evaluation and testing?
Advertisement
Answer
def load_checkpoint(checkpoint, model, optimizer = None): if not os.path.exists(checkpoint): raise("File does not exists {}".format(checkpoint)) checkpoint = torch.load(checkpoint) model.load_state_dict(checkpoint['state_dict']) if optimizer: optimizer.load_state_dict(checkpoint['optim_dict']) return checkpoint
To test a model you need to load the state dictionary of your trained model and optimizer (if applicable). But, if you are resuming training from a point and you are using any sort of scheduler you need to load the scheduler state too.