Model saving is taking a lot of time when using dicts

save = {}
save["model_weights"] = self.model.state_dict()
save["epochs_completed"] = ep
save["optimizer"]=self.optimizer.state_dict()
save["best_score"] = score
pickle.dump(save, open(self.opt.save_loc+self.opt.save_name, "wb"))

The above save method is taking a lot of time. I am basically doing this to resume my training. Is there any mistake I am doing ?

pickle isn’t optimised for saving huge arrays of data. Try using torch.save instead.

1 Like