def train_model():
train_dataset = torch.load(
os.path.join(paths.experiment_data_dir, 'train.data'))
dev_dataset = torch.load(
os.path.join(paths.experiment_data_dir, 'dev.data'))
test_dataset = torch.load(
os.path.join(paths.experiment_data_dir, 'test.data'))
vocab = torch.load(paths.vocab_path)
model = RelationClassifier(encoder_type=args.encoder,
num_classes=len(train_dataset.label_map),
vocab_size=vocab.size(),
embed_dim=vocab.embed_dim,
model_config=ModelConfig,
drop_rate=args.drop,
use_cuda=args.cuda,
attention=args.attention)
# init with pre-trained embeddings
model.argument_encoder.emb.weight.data.copy_(vocab.embeddings)
# model.argument_encoder.emb.weight.requires_grad = False
criterion = nn.CrossEntropyLoss()
params_to_train = filter(lambda p: p.requires_grad, model.parameters())
if args.cuda:
model.cuda(), criterion.cuda()
if args.optim == 'sgd':
optimizer = optim.SGD(params_to_train,
lr=args.lr,
weight_decay=args.wd)
elif args.optim == 'adam':
optimizer = optim.Adam(params_to_train,
lr=args.lr,
weight_decay=args.wd)
elif args.optim == 'adagrad':
optimizer = optim.Adagrad(params_to_train,
lr=args.lr,
weight_decay=args.wd)
elif args.optim == 'rprop':
optimizer = optim.Rprop(params_to_train, lr=args.lr)
trainer = Trainer(model, criterion, optimizer, args.cuda)
best_dev_acc = 0
for epoch in range(args.epochs):
print()
trainer.train(train_dataset, vocab, batch_size=args.batch_size)
# train_loss, train_acc, train_f1 = trainer.eval(train_dataset, vocab, 'trainset')
# print('Train: loss {}, acc {}, f1 {}'.format(train_loss, train_acc, train_f1))
dev_acc = trainer.eval(dev_dataset, vocab, 'devset')
print('Dev: acc {}'.format(dev_acc))
test_acc = trainer.eval(test_dataset, vocab, 'testset')
print('Test: acc {}'.format(test_acc))
if dev_acc > best_dev_acc:
best_dev_acc = dev_acc
print('Model saved to {}'.format(paths.best_model_path))
torch.save(model.state_dict(), paths.best_model_path)
def test_model():
train_dataset = torch.load(
os.path.join(paths.experiment_data_dir, 'train.data'))
dev_dataset = torch.load(
os.path.join(paths.experiment_data_dir, 'dev.data'))
test_dataset = torch.load(
os.path.join(paths.experiment_data_dir, 'test.data'))
vocab = torch.load(paths.vocab_path)
model = RelationClassifier(encoder_type=args.encoder,
num_classes=len(train_dataset.label_map),
vocab_size=vocab.size(),
embed_dim=vocab.embed_dim,
model_config=ModelConfig,
drop_rate=args.drop,
use_cuda=args.cuda)
model.load_state_dict(torch.load(paths.best_model_path))
trainer = Trainer(model, use_cuda=args.cuda)
# train_loss, train_acc, train_f1 = trainer.eval(train_dataset, vocab, 'trainset')
# print('Train: loss {}, acc {}, f1 {}'.format(train_loss, train_acc, train_f1))
dev_acc = trainer.eval(dev_dataset, vocab, 'devset')
print('Dev: acc {}'.format(dev_acc))
test_acc = trainer.eval(test_dataset, vocab, 'testset')
print('Test: acc {}'.format(test_acc))
