def main():
opt = parser.parse_args()
checkpoint = torch.load(opt.model)
opt.cuda = opt.gpu > -1
if opt.cuda:
torch.cuda.set_device(opt.gpu)
model_opt = checkpoint['opt']
src_dict = checkpoint['dicts']['src']
tgt_dict = checkpoint['dicts']['tgt']
feature_dicts = []
embeddings = make_embeddings(model_opt, src_dict, feature_dicts)
encoder = make_encoder(model_opt, embeddings)
embeddings = make_embeddings(model_opt,
tgt_dict,
feature_dicts,
for_encoder=False)
decoder = make_decoder(model_opt, embeddings)
encoder_embeddings = encoder.word_lut.weight.data.tolist()
decoder_embeddings = decoder.word_lut.weight.data.tolist()
print("Writing source embeddings")
write_embeddings(opt.output_dir + "/src_embeddings.txt", src_dict,
encoder_embeddings)
print("Writing target embeddings")
write_embeddings(opt.output_dir + "/tgt_embeddings.txt", tgt_dict,
decoder_embeddings)
print('... done.')
print('Converting model...')
def nmtmodel_forward(self, opt, source_l=3, bsize=1):
"""
Creates a nmtmodel with a custom opt function.
Forwards a testbatch and checks output size.
Args:
opt: Namespace with options
source_l: length of input sequence
bsize: batchsize
"""
word_dict = self.get_vocab()
feature_dicts = []
embeddings = make_embeddings(opt, word_dict, feature_dicts)
enc = make_encoder(opt, embeddings)
embeddings = make_embeddings(opt, word_dict, feature_dicts,
for_encoder=False)
dec = make_decoder(opt, embeddings)
model = onmt.Models.NMTModel(enc, dec)
test_src, test_tgt, test_length = self.get_batch(source_l=source_l,
bsize=bsize)
outputs, attn, _ = model(test_src,
test_tgt,
test_length)
outputsize = torch.zeros(source_l - 1, bsize, opt.rnn_size)
# Make sure that output has the correct size and type
self.assertEqual(outputs.size(), outputsize.size())
self.assertEqual(type(outputs), torch.autograd.Variable)
self.assertEqual(type(outputs.data), torch.FloatTensor)
def encoder_forward(self, opt, source_l=3, bsize=1):
'''
Tests if the encoder works as expected
args:
opt: set of options
source_l: Length of generated input sentence
bsize: Batchsize of generated input
'''
word_dict = self.get_vocab()
feature_dicts = []
embeddings = make_embeddings(opt, word_dict, feature_dicts)
enc = make_encoder(opt, embeddings)
test_src, test_tgt, test_length = self.get_batch(source_l=source_l,
bsize=bsize)
hidden_t, outputs = enc(test_src, test_length)
# Initialize vectors to compare size with
test_hid = torch.zeros(self.opt.enc_layers, bsize, opt.rnn_size)
test_out = torch.zeros(source_l, bsize, opt.rnn_size)
# Ensure correct sizes and types
self.assertEqual(test_hid.size(),
hidden_t[0].size(),
hidden_t[1].size())
self.assertEqual(test_out.size(), outputs.size())
self.assertEqual(type(outputs), torch.autograd.Variable)
self.assertEqual(type(outputs.data), torch.FloatTensor)
def encoder_forward(self, opt, source_l=3, bsize=1):
'''
Tests if the encoder works as expected
args:
opt: set of options
source_l: Length of generated input sentence
bsize: Batchsize of generated input
'''
word_dict = self.get_vocab()
feature_dicts = []
embeddings = make_embeddings(opt, word_dict, feature_dicts)
enc = make_encoder(opt, embeddings)
test_src, test_tgt, test_length = self.get_batch(source_l=source_l,
bsize=bsize)
hidden_t, outputs = enc(test_src, test_length)
# Initialize vectors to compare size with
test_hid = torch.zeros(self.opt.enc_layers, bsize, opt.rnn_size)
test_out = torch.zeros(source_l, bsize, opt.rnn_size)
# Ensure correct sizes and types
self.assertEqual(test_hid.size(), hidden_t[0].size(),
hidden_t[1].size())
self.assertEqual(test_out.size(), outputs.size())
self.assertEqual(type(outputs), torch.autograd.Variable)
self.assertEqual(type(outputs.data), torch.FloatTensor)
def nmtmodel_forward(self, opt, source_l=3, bsize=1):
"""
Creates a nmtmodel with a custom opt function.
Forwards a testbatch and checks output size.
Args:
opt: Namespace with options
source_l: length of input sequence
bsize: batchsize
"""
word_dict = self.get_vocab()
feature_dicts = []
embeddings = make_embeddings(opt, word_dict, feature_dicts)
enc = make_encoder(opt, embeddings)
embeddings = make_embeddings(opt,
word_dict,
feature_dicts,
for_encoder=False)
dec = make_decoder(opt, embeddings)
model = onmt.Models.NMTModel(enc, dec)
test_src, test_tgt, test_length = self.get_batch(source_l=source_l,
bsize=bsize)
outputs, attn, _ = model(test_src, test_tgt, test_length)
outputsize = torch.zeros(source_l - 1, bsize, opt.rnn_size)
# Make sure that output has the correct size and type
self.assertEqual(outputs.size(), outputsize.size())
self.assertEqual(type(outputs), torch.autograd.Variable)
self.assertEqual(type(outputs.data), torch.FloatTensor)
def main():
opt = parser.parse_args()
checkpoint = torch.load(opt.model,
map_location=lambda storage, loc: storage)
#checkpoint = torch.load('../200k-model_acc_47.70_ppl_20.36_e13.pt',map_location=lambda storage, loc: storage)
opt.cuda = opt.gpu > -1
if opt.cuda:
torch.cuda.set_device(opt.gpu)
fields = onmt.IO.load_fields(checkpoint['vocab'])
src_dict = fields["src"].vocab
tgt_dict = fields["tgt"].vocab
feature_dicts = onmt.IO.collect_feature_dicts(fields, 'src')
model_opt = checkpoint['opt']
# src_dict = checkpoint['dicts']['src']
# tgt_dict = checkpoint['dicts']['tgt']
# feature_dicts = []
embeddings = make_embeddings(model_opt, src_dict, feature_dicts)
encoder = make_encoder(model_opt, embeddings)
embeddings = make_embeddings(model_opt,
tgt_dict,
feature_dicts,
for_encoder=False)
decoder = make_decoder(model_opt, embeddings)
encoder_embeddings = encoder.embeddings.word_lut.weight.data.tolist()
decoder_embeddings = decoder.embeddings.word_lut.weight.data.tolist()
print("Writing source embeddings")
write_embeddings(opt.output_dir + "/src_embeddings.txt", src_dict,
encoder_embeddings)
print("Writing target embeddings")
write_embeddings(opt.output_dir + "/tgt_embeddings.txt", tgt_dict,
decoder_embeddings)
print('... done.')
print('Converting model...')
def ntmmodel_forward(self, opt, sourceL=3, bsize=1):
"""
Creates a ntmmodel with a custom opt function.
Forwards a testbatch anc checks output size.
Args:
opt: Namespace with options
sourceL: length of input sequence
bsize: batchsize
"""
vocab = self.get_vocab()
word_padding_idx = vocab.stoi[onmt.IO.PAD_WORD]
feats_padding_idx = []
embeddings = make_embeddings(opt,
word_padding_idx,
feats_padding_idx,
len(vocab),
for_encoder=True)
enc = make_encoder(opt, embeddings)
embeddings = make_embeddings(opt,
word_padding_idx,
feats_padding_idx,
len(vocab),
for_encoder=False)
dec = make_decoder(opt, embeddings)
model = onmt.Models.NMTModel(enc, dec)
test_src, test_tgt, test_length = self.get_batch(sourceL=sourceL,
bsize=bsize)
outputs, attn, _ = model(test_src, test_tgt, test_length)
outputsize = torch.zeros(sourceL - 1, bsize, opt.rnn_size)
# Make sure that output has the correct size and type
self.assertEqual(outputs.size(), outputsize.size())
self.assertEqual(type(outputs), torch.autograd.Variable)
self.assertEqual(type(outputs.data), torch.FloatTensor)
def make_base_model(model_opt, fields, gpu, checkpoint=None):
"""
Args:
model_opt: the option loaded from checkpoint.
fields: `Field` objects for the model.
gpu(bool): whether to use gpu.
checkpoint: the model gnerated by train phase, or a resumed snapshot
model from a stopped training.
Returns:
the NMTModel.
"""
assert model_opt.model_type in ["text", "img", "audio"], \
("Unsupported model type %s" % (model_opt.model_type))
# Make encoder.
if model_opt.model_type == "text":
src_dict = fields["src"].vocab
feature_dicts = onmt.io.collect_feature_vocabs(fields, 'src')
src_embeddings = make_embeddings(model_opt, src_dict, feature_dicts)
encoder = make_encoder(model_opt, src_embeddings)
elif model_opt.model_type == "img":
encoder = ImageEncoder(model_opt.enc_layers, model_opt.brnn,
model_opt.rnn_size, model_opt.dropout)
elif model_opt.model_type == "audio":
encoder = AudioEncoder(model_opt.enc_layers, model_opt.brnn,
model_opt.rnn_size, model_opt.dropout,
model_opt.sample_rate, model_opt.window_size)
# Make decoder.
tgt_dict = fields["tgt"].vocab
feature_dicts = onmt.io.collect_feature_vocabs(fields, 'tgt')
tgt_embeddings = make_embeddings(model_opt,
tgt_dict,
feature_dicts,
for_encoder=False)
# Share the embedding matrix - preprocess with share_vocab required.
if model_opt.share_embeddings:
# src/tgt vocab should be the same if `-share_vocab` is specified.
if src_dict != tgt_dict:
raise AssertionError('The `-share_vocab` should be set during '
'preprocess if you use share_embeddings!')
tgt_embeddings.word_lut.weight = src_embeddings.word_lut.weight
decoder = make_decoder(model_opt, tgt_embeddings)
# Make NMTModel(= encoder + decoder).
model = NMTTemplateModel(encoder, decoder, model_opt.rnn_size)
model.model_type = model_opt.model_type
# Make Generator.
if not model_opt.copy_attn:
generator = nn.Sequential(
nn.Linear(model_opt.rnn_size, len(fields["tgt"].vocab)),
nn.LogSoftmax())
if model_opt.share_decoder_embeddings:
generator[0].weight = decoder.embeddings.word_lut.weight
else:
generator = CopyGenerator(model_opt.rnn_size, fields["tgt"].vocab)
# Load the model states from checkpoint or initialize them.
if checkpoint is not None:
print('Loading model parameters.')
model.load_state_dict(checkpoint['model'])
generator.load_state_dict(checkpoint['generator'])
else:
if model_opt.param_init != 0.0:
print('Intializing model parameters.')
for p in model.parameters():
p.data.uniform_(-model_opt.param_init, model_opt.param_init)
for p in generator.parameters():
p.data.uniform_(-model_opt.param_init, model_opt.param_init)
if hasattr(model.encoder, 'embeddings'):
model.encoder.embeddings.load_pretrained_vectors(
model_opt.pre_word_vecs_enc, model_opt.fix_word_vecs_enc)
if hasattr(model.decoder, 'embeddings'):
model.decoder.embeddings.load_pretrained_vectors(
model_opt.pre_word_vecs_dec, model_opt.fix_word_vecs_dec)
# Add generator to model (this registers it as parameter of model).
model.generator = generator
# Make the whole model leverage GPU if indicated to do so.
if gpu:
model.cuda()
else:
model.cpu()
return model
