def load_model(args): rnn = models.DecoderRNN() model_dir_rnn = args.resume model_std = torch.load(os.path.join(model_dir_rnn, 'model_best_rnn.pth.tar'), map_location="cuda:"+str(args.gpu)) rnn.load_state_dict(model_std) rnn = rnn.cuda() feature_stractor = models.Stractor() feature_stractor = feature_stractor.cuda() #model = torch.nn.DataParallel(model, #device_ids=list(range(torch.cuda.device_count()))).cuda() return feature_stractor, rnn
def load_model(args): classifier = models.Classifier() feature = models.Stractor() model_dir_class = os.path.join(args.resume) model_std = torch.load(os.path.join(model_dir_class, 'model_best_class.pth.tar'), map_location="cuda:" + str(args.gpu)) classifier.load_state_dict(model_std) classifier = classifier.cuda() feature = feature.cuda() #model = torch.nn.DataParallel(model, #device_ids=list(range(torch.cuda.device_count()))).cuda() return feature, classifier
def load_model(args): classifier = models.Classifier() stract = models.Stractor() model_dir_class = os.path.join(args.resume, 'Classifier') model_dir_str = os.path.join(args.resume, 'featureStractor') model_std = torch.load(os.path.join(model_dir_class, 'model_1_class.pth.tar'), map_location="cuda:"+str(args.gpu)) classifier.load_state_dict(model_std) classifier = classifier.cuda() model_std = torch.load(os.path.join(model_dir_str, 'model_1_feaStr.pth.tar'), map_location="cuda:"+str(args.gpu)) stract.load_state_dict(model_std) stract = stract.cuda() #model = torch.nn.DataParallel(model, #device_ids=list(range(torch.cuda.device_count()))).cuda() return stract, classifier
np.random.seed(args.random_seed) torch.manual_seed(args.random_seed) torch.cuda.manual_seed(args.random_seed) ''' load dataset and prepare data loader ''' print('===> prepare dataloader ...') train_loader = torch.utils.data.DataLoader(data.DATA(args, mode='train'), batch_size=args.train_batch, num_workers=args.workers, shuffle=True) val_loader = torch.utils.data.DataLoader(data.DATA(args, mode='valid'), batch_size=args.train_batch, num_workers=args.workers, shuffle=False) ''' load model ''' print('===> prepare model ...') feature_stractor = models.Stractor() feature_stractor.cuda() # load model to gpu params_to_update = feature_stractor.parameters() params_to_update_str = [] for name, param in feature_stractor.named_parameters(): if param.requires_grad == True: params_to_update_str.append(param) classifier = models.Classifier() classifier = classifier.cuda() params_to_update_class = [] for name, param in classifier.named_parameters(): if param.requires_grad == True: params_to_update_class.append(param) ''' define loss ''' criterion = nn.CrossEntropyLoss()