def build_model(opt, vocab, checkpoint, gpu_id=None):
print('Building model...')
use_gpu = len(opt.gpu_ranks) > 0
if use_gpu and gpu_id is not None:
device = torch.device("cuda", gpu_id)
elif use_gpu and not gpu_id:
device = torch.device("cuda")
elif not use_gpu:
device = torch.device("cpu")
# model = MLP([1000, 500, 500], len(vocab))
model = CNN(len(vocab), opt.dropout)
if checkpoint is not None:
# end of patch for backward compatibility
if opt.train_from:
print("Loading model parameters from checkpoint...")
model.load_state_dict(checkpoint['model'], strict=False)
else:
print("Loading model parameters from checkpoint...")
cur_state_dict = model.state_dict()
for key in cur_state_dict.keys():
if key in checkpoint['model'] and cur_state_dict[key].size(
) == checkpoint['model'][key].size():
cur_state_dict[key] = checkpoint['model'][key]
# elif key in checkpoint['model'] and cur_state_dict[key].size() != checkpoint['model'][key].size():
# print("***" , key)
model.load_state_dict(cur_state_dict, strict=False)
if len(opt.gpu_ranks) > 1:
model = nn.DataParallel(model, opt.gpu_ranks)
model.to(device)
return model, device
def build_model(opt, vocab, checkpoint, gpu_id=None):
print('Building model...')
if opt.gpu == -1:
device = torch.device("cpu")
else:
device = torch.device("cuda", gpu_id)
# model = MLP([1000, 500, 500], len(vocab))
model = CNN(len(vocab))
if checkpoint is not None:
# end of patch for backward compatibility
print("Loading model parameters from checkpoint...")
model.load_state_dict(checkpoint['model'], strict=False)
model.to(device)
return model, device
def load_model(opt):
global model
global device
global idx2label
print('Loading checkpoint from %s' % opt.model)
checkpoint = torch.load(opt.model,
map_location=lambda storage, loc: storage)
print('Loading vocab from checkpoint at %s' % opt.model)
vocab = checkpoint['vocab']
idx2label = {v: k for k, v in vocab.items()}
print('Building model...')
if opt.gpu == -1:
device = torch.device("cpu")
else:
device = torch.device("cuda", opt.gpu)
model = CNN(len(vocab))
# end of patch for backward compatibility
print("Loading model parameters from checkpoint...")
model.load_state_dict(checkpoint['model'], strict=False)
model.to(device)
model.eval()