def load_pretrained(self, model_folder, iteration=None): """ Loads network and optimizer parameters from a previous training run. Note: It is up to the user to ensure that the model parameters match. e.g. width, height, descriptor dimension etc. :param model_folder: location of the folder containing the param files 001000.pth. Can be absolute or relative path. If relative then it is relative to pdc/trained_models/ :type model_folder: :param iteration: which index to use, e.g. 3500, if None it loads the latest one :type iteration: :return: iteration :rtype: """ if not os.path.isdir(model_folder): pdc_path = utils.getPdcPath() model_folder = os.path.join(pdc_path, "trained_models", model_folder) # find idx.pth and idx.pth.opt files if iteration is None: files = os.listdir(model_folder) model_param_file = sorted(fnmatch.filter(files, '*.pth'))[-1] iteration = int(model_param_file.split(".")[0]) optim_param_file = sorted(fnmatch.filter(files, '*.pth.opt'))[-1] else: prefix = utils.getPaddedString(iteration, width=6) model_param_file = prefix + ".pth" optim_param_file = prefix + ".pth.opt" print("model_param_file", model_param_file) model_param_file = os.path.join(model_folder, model_param_file) optim_param_file = os.path.join(model_folder, optim_param_file) self._dcn = self.build_network() self._dcn.load_state_dict(torch.load(model_param_file)) self._dcn.cuda() self._dcn.train() self._optimizer = self._construct_optimizer(self._dcn.parameters()) self._optimizer.load_state_dict(torch.load(optim_param_file)) return iteration
def load_pretrained(self, model_folder, iteration=None): """ Loads network and optimizer parameters from a previous training run. Note: It is up to the user to ensure that the model parameters match. e.g. width, height, descriptor dimension etc. :param model_folder: location of the folder containing the param files 001000.pth. Can be absolute or relative path. If relative then it is relative to pdc/trained_models/ :type model_folder: :param iteration: which index to use, e.g. 3500, if None it loads the latest one :type iteration: :return: iteration :rtype: """ if not os.path.isdir(model_folder): pdc_path = utils.getPdcPath() model_folder = os.path.join(pdc_path, "trained_models", model_folder) # find idx.pth and idx.pth.opt files if iteration is None: files = os.listdir(model_folder) model_param_file = sorted(fnmatch.filter(files, '*.pth'))[-1] iteration = int(model_param_file.split(".")[0]) optim_param_file = sorted(fnmatch.filter(files, '*.pth.opt'))[-1] else: prefix = utils.getPaddedString(iteration, width=6) model_param_file = prefix + ".pth" optim_param_file = prefix + ".pth.opt" print "model_param_file", model_param_file model_param_file = os.path.join(model_folder, model_param_file) optim_param_file = os.path.join(model_folder, optim_param_file) self._dcn = self.build_network() self._dcn.load_state_dict(torch.load(model_param_file)) self._dcn.cuda() self._dcn.train() self._optimizer = self._construct_optimizer(self._dcn.parameters()) self._optimizer.load_state_dict(torch.load(optim_param_file)) return iteration