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()
