def run(args):
writer = SummaryWriter()
src, tgt, _, _ = build_dataset(args)
print('Loading test data split.')
_, _, test_gen = datasets.Multi30k.splits(
exts=(build_file_extension(args.src_language),
build_file_extension(args.tgt_language)),
fields=(('src', src), ('tgt', tgt)),
filter_pred=lambda x: len(vars(x)['src']) <= args.max_seq_length and
len(vars(x)['tgt']) <= args.max_seq_length)
print('Finished loading test data split.')
src_vocab_size = len(src.vocab.itos)
tgt_vocab_size = len(tgt.vocab.itos)
_, _, test_iterator = data.Iterator.splits(
(_, _, test_gen),
sort_key=lambda x: len(x.src),
batch_sizes=(args.batch_size, args.batch_size, args.batch_size))
print('Instantiating model...')
device = args.device
model = Transformer(src_vocab_size,
tgt_vocab_size,
device,
p_dropout=args.dropout)
model = model.to(device)
model.load_state_dict(torch.load(args.model))
print('Model instantiated!')
print('Starting testing...')
test(model, test_iterator, src.vocab, tgt.vocab, args, writer)
print('Finished testing.')
def __init__(self, src_vocab, tgt_vocab, src_vocab_size, tgt_vocab_size,
args):
self.max_seq_length = args.max_seq_length
self.device = args.device
self.src_vocab = src_vocab
self.tgt_vocab = tgt_vocab
self.beam_size = args.beam_size
model = Transformer(src_vocab_size, tgt_vocab_size, args.device)
model.load_state_dict(torch.load(args.model))
model = model.to(args.device)
self.model = model
self.model.eval()
def run(args):
writer = SummaryWriter()
src, tgt, train_iterator, val_iterator = build_dataset(args)
src_vocab_size = len(src.vocab.itos)
tgt_vocab_size = len(tgt.vocab.itos)
print('Instantiating model...')
device = args.device
model = Transformer(src_vocab_size,
tgt_vocab_size,
device,
p_dropout=args.dropout)
model = model.to(device)
if args.checkpoint is not None:
model.load_state_dict(torch.load(args.checkpoint))
else:
for p in model.parameters():
if p.dim() > 1:
nn.init.xavier_uniform_(p)
print('Model instantiated!')
optimizer = optim.Adam(model.parameters(),
lr=args.lr,
betas=(0.9, 0.98),
eps=1e-9)
print('Starting training...')
for epoch in range(args.epochs):
acc = train(model, epoch + 1, train_iterator, optimizer, src.vocab,
tgt.vocab, args, writer)
model_file = 'models/model_' + str(epoch) + '_' + str(acc) + '.pth'
torch.save(model.state_dict(), model_file)
print('Saved model to ' + model_file)
validate(model, epoch + 1, val_iterator, src.vocab, tgt.vocab, args,
writer)
print('Finished training.')
import numpy as np
import torch
from config import device, logger, data_file, vocab_file
from transformer.transformer import Transformer
if __name__ == '__main__':
# filename = 'transformer.pt'
filename = 'BEST'
print('loading {}...'.format(filename))
start = time.time()
model = Transformer()
model.load_state_dict(torch.load(filename))
print('elapsed {} sec'.format(time.time() - start))
model = model.to(device)
model.eval()
assert (1 == 0)
logger.info('loading samples...')
start = time.time()
with open(data_file, 'rb') as file:
data = pickle.load(file)
samples = data['valid']
elapsed = time.time() - start
logger.info('elapsed: {:.4f} seconds'.format(elapsed))
logger.info('loading vocab...')
start = time.time()
with open(vocab_file, 'rb') as file:
data = pickle.load(file)
src_idx2char = data['dict']['src_idx2char']
def train_net(args):
torch.manual_seed(7)
np.random.seed(7)
checkpoint = args.checkpoint
start_epoch = 0
best_loss = float('inf')
writer = SummaryWriter()
epochs_since_improvement = 0
# Initialize / load checkpoint
if checkpoint is None:
# model
encoder = Encoder(n_src_vocab,
args.n_layers_enc,
args.n_head,
args.d_k,
args.d_v,
args.d_model,
args.d_inner,
dropout=args.dropout,
pe_maxlen=args.pe_maxlen)
decoder = Decoder(
sos_id,
eos_id,
n_tgt_vocab,
args.d_word_vec,
args.n_layers_dec,
args.n_head,
args.d_k,
args.d_v,
args.d_model,
args.d_inner,
dropout=args.dropout,
tgt_emb_prj_weight_sharing=args.tgt_emb_prj_weight_sharing,
pe_maxlen=args.pe_maxlen)
model = Transformer(encoder, decoder)
# print(model)
# model = nn.DataParallel(model)
# optimizer
optimizer = TransformerOptimizer(
torch.optim.Adam(model.parameters(), betas=(0.9, 0.98), eps=1e-09))
else:
checkpoint = torch.load(checkpoint)
start_epoch = checkpoint['epoch'] + 1
epochs_since_improvement = checkpoint['epochs_since_improvement']
model = checkpoint['model']
optimizer = checkpoint['optimizer']
# Move to GPU, if available
model = model.to(device)
# Custom dataloaders
train_dataset = AiChallenger2017Dataset('train')
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
collate_fn=pad_collate,
shuffle=True,
num_workers=args.num_workers)
valid_dataset = AiChallenger2017Dataset('valid')
valid_loader = torch.utils.data.DataLoader(valid_dataset,
batch_size=args.batch_size,
collate_fn=pad_collate,
shuffle=False,
num_workers=args.num_workers)
# Epochs
for epoch in range(start_epoch, args.epochs):
# One epoch's training
train_loss = train(train_loader=train_loader,
model=model,
optimizer=optimizer,
epoch=epoch,
logger=logger,
writer=writer)
writer.add_scalar('epoch/train_loss', train_loss, epoch)
writer.add_scalar('epoch/learning_rate', optimizer.lr, epoch)
print('\nLearning rate: {}'.format(optimizer.lr))
print('Step num: {}\n'.format(optimizer.step_num))
# One epoch's validation
valid_loss = valid(valid_loader=valid_loader,
model=model,
logger=logger)
writer.add_scalar('epoch/valid_loss', valid_loss, epoch)
# Check if there was an improvement
is_best = valid_loss < best_loss
best_loss = min(valid_loss, best_loss)
if not is_best:
epochs_since_improvement += 1
print("\nEpochs since last improvement: %d\n" %
(epochs_since_improvement, ))
else:
epochs_since_improvement = 0
# Save checkpoint
save_checkpoint(epoch, epochs_since_improvement, model, optimizer,
best_loss, is_best)
def train_net(args):
# 为了保证程序执行结果一致, 给随机化设定种子
torch.manual_seed(7)
np.random.seed(7)
checkpoint = args.checkpoint
start_epoch = 0
writer = SummaryWriter()
if checkpoint is None:
# model
encoder = Encoder(Config.vocab_size, args.n_layers_enc, args.n_head,
args.d_k, args.d_v, args.d_model, args.d_inner,
dropout=args.dropout, pe_maxlen=args.pe_maxlen)
decoder = Decoder(Config.sos_id, Config.eos_id, Config.vocab_size,
args.d_word_vec, args.n_layers_dec, args.n_head,
args.d_k, args.d_v, args.d_model, args.d_inner,
dropout=args.dropout,
tgt_emb_prj_weight_sharing=args.tgt_emb_prj_weight_sharing,
pe_maxlen=args.pe_maxlen)
model = Transformer(encoder, decoder)
# optimizer
optimizer = TransformerOptimizer(
torch.optim.Adam(model.parameters(), betas=(0.9, 0.98), eps=1e-09))
else:
checkpoint = torch.load(checkpoint)
start_epoch = checkpoint['epoch'] + 1
model = checkpoint['model']
optimizer = checkpoint['optimizer']
# Move to GPU, if available
model = model.to(Config.device)
# Custom dataloaders 数据的加载 注意这里指定了一个参数collate_fn代表的数据需要padding
train_dataset = TranslateDataset()
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, collate_fn=pad_collate,
shuffle=True, num_workers=args.num_workers)
# Epochs
Loss_list = []
for epoch in range(start_epoch, args.epochs):
# One epoch's training
train_loss = train(train_loader=train_loader,
model=model,
optimizer=optimizer,
epoch=epoch,
logger=logger,
writer=writer)
l = str(train_loss)
Loss_list.append(l)
l_temp = l + '\n'
with open('loss_epoch.txt', 'a+') as f:
f.write(l_temp)
writer.add_scalar('epoch/train_loss', train_loss, epoch)
writer.add_scalar('epoch/learning_rate', optimizer.lr, epoch)
print('\nLearning rate: {}'.format(optimizer.lr))
print('Step num: {}\n'.format(optimizer.step_num))
# Save checkpoint
save_checkpoint(epoch, model, optimizer, train_loss)
with open('loss.txt', 'w') as f:
f.write('\n'.join(Loss_list))
from config import Config, logger
from transformer.transformer import Transformer
if __name__ == '__main__':
# 先去执行export.py 把模型导出来
filename = 'reading_comprehension.pt' # 导出模型所放的位置
print('loading {}...'.format(filename))
start = time.time()
model = Transformer()
model.load_state_dict(torch.load(filename))
print('elapsed {} sec'.format(time.time() - start))
model = model.to(Config.device)
model.eval()
# 加载测试集
logger.info('loading samples...')
start = time.time()
with open(Config.data_file, 'rb') as file:
data = pickle.load(file)
samples = data
elapsed = time.time() - start
logger.info('elapsed: {:.4f} seconds'.format(elapsed))
# 加载词典
logger.info('loading vocab...')
start = time.time()
with open(Config.vocab_file, 'rb') as file:
