def save_to_file(self, file_name): """Save this task to file.""" # Save trajectories from new epochs first. epochs_to_save = [e for e in self._trajectories if e not in self._saved_epochs_unchanged] for epoch in epochs_to_save: training.pickle_to_file(self._trajectories[epoch], self._epoch_filename(file_name, epoch), gzip=True) # Now save the list of epochs (so the trajectories are already there, # even in case of preemption). dictionary = {'n_interactions': self._n_interactions, 'n_trajectories': self._n_trajectories, 'max_steps': self._max_steps, 'gamma': self._gamma, 'all_epochs': list(self._trajectories.keys())} training.pickle_to_file(dictionary, file_name, gzip=False)
def _save_state_dict(self, trainer_state_dict, weights_file): training.pickle_to_file(trainer_state_dict, weights_file, gzip=True) log('Model saved to %s' % weights_file, stdout=False)