mirror of
https://github.com/wassname/Castor.git
synced 2026-09-09 11:13:20 +08:00
211 lines
8.1 KiB
Python
211 lines
8.1 KiB
Python
import time
|
|
import os
|
|
import numpy as np
|
|
import random
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.onnx
|
|
from torchtext import data
|
|
|
|
from args import get_args
|
|
from model import SmPlusPlus
|
|
from utils.relevancy_metrics import get_map_mrr
|
|
from trec_dataset import TrecDataset
|
|
from wiki_dataset import WikiDataset
|
|
|
|
args = get_args()
|
|
config = args
|
|
|
|
torch.manual_seed(args.seed)
|
|
|
|
def set_vectors(field, vector_path):
|
|
if os.path.isfile(vector_path):
|
|
stoi, vectors, dim = torch.load(vector_path)
|
|
field.vocab.vectors = torch.Tensor(len(field.vocab), dim)
|
|
|
|
for i, token in enumerate(field.vocab.itos):
|
|
wv_index = stoi.get(token, None)
|
|
if wv_index is not None:
|
|
field.vocab.vectors[i] = vectors[wv_index]
|
|
else:
|
|
# initialize <unk> with U(-0.25, 0.25) vectors
|
|
field.vocab.vectors[i] = torch.FloatTensor(dim).uniform_(-0.25, 0.25)
|
|
else:
|
|
print("Error: Need word embedding pt file")
|
|
exit(1)
|
|
return field
|
|
|
|
# Set default configuration in : args.py
|
|
args = get_args()
|
|
config = args
|
|
|
|
# Set random seed for reproducibility
|
|
torch.manual_seed(args.seed)
|
|
torch.backends.cudnn.deterministic = True
|
|
if not args.cuda:
|
|
args.gpu = -1
|
|
if torch.cuda.is_available() and args.cuda:
|
|
print("Note: You are using GPU for training")
|
|
torch.cuda.set_device(args.gpu)
|
|
torch.cuda.manual_seed(args.seed)
|
|
if torch.cuda.is_available() and not args.cuda:
|
|
print("You have Cuda but you're using CPU for training.")
|
|
np.random.seed(args.seed)
|
|
random.seed(args.seed)
|
|
|
|
QID = data.Field(sequential=False)
|
|
QUESTION = data.Field(batch_first=True)
|
|
ANSWER = data.Field(batch_first=True)
|
|
LABEL = data.Field(sequential=False)
|
|
EXTERNAL = data.Field(sequential=True, dtype=torch.FloatTensor, batch_first=True, use_vocab=False,
|
|
postprocessing=data.Pipeline(lambda arr, _, train: [float(y) for y in arr]))
|
|
|
|
if config.dataset == 'TREC':
|
|
train, dev, test = TrecDataset.splits(QID, QUESTION, ANSWER, EXTERNAL, LABEL)
|
|
elif config.dataset == 'wiki':
|
|
train, dev, test = WikiDataset.splits(QID, QUESTION, ANSWER, EXTERNAL, LABEL)
|
|
else:
|
|
print("Unsupported dataset")
|
|
exit()
|
|
|
|
QID.build_vocab(train, dev, test)
|
|
QUESTION.build_vocab(train, dev, test)
|
|
ANSWER.build_vocab(train, dev, test)
|
|
LABEL.build_vocab(train, dev, test)
|
|
|
|
|
|
QUESTION = set_vectors(QUESTION, args.vector_cache)
|
|
ANSWER = set_vectors(ANSWER, args.vector_cache)
|
|
|
|
train_iter = data.Iterator(train, batch_size=args.batch_size, device=args.gpu, train=True, repeat=False,
|
|
sort=False, shuffle=True)
|
|
dev_iter = data.Iterator(dev, batch_size=args.batch_size, device=args.gpu, train=False, repeat=False,
|
|
sort=False, shuffle=False)
|
|
test_iter = data.Iterator(test, batch_size=args.batch_size, device=args.gpu, train=False, repeat=False,
|
|
sort=False, shuffle=False)
|
|
|
|
config.target_class = len(LABEL.vocab)
|
|
config.questions_num = len(QUESTION.vocab)
|
|
config.answers_num = len(ANSWER.vocab)
|
|
|
|
print("Dataset {} Mode {}".format(args.dataset, args.mode))
|
|
print("VOCAB num", len(QUESTION.vocab))
|
|
print("LABEL.target_class:", len(LABEL.vocab))
|
|
print("LABELS:", LABEL.vocab.itos)
|
|
print("Train instance", len(train))
|
|
print("Dev instance", len(dev))
|
|
print("Test instance", len(test))
|
|
|
|
if args.resume_snapshot:
|
|
if args.cuda:
|
|
model = torch.load(args.resume_snapshot, map_location=lambda storage, location: storage.cuda(args.gpu))
|
|
else:
|
|
model = torch.load(args.resume_snapshot, map_location=lambda storage, location: storage)
|
|
else:
|
|
model = SmPlusPlus(config)
|
|
model.static_question_embed.weight.data.copy_(QUESTION.vocab.vectors)
|
|
model.nonstatic_question_embed.weight.data.copy_(QUESTION.vocab.vectors)
|
|
model.static_answer_embed.weight.data.copy_(ANSWER.vocab.vectors)
|
|
model.nonstatic_answer_embed.weight.data.copy_(ANSWER.vocab.vectors)
|
|
|
|
if args.cuda:
|
|
model.cuda()
|
|
print("Shift model to GPU")
|
|
|
|
|
|
parameter = filter(lambda p: p.requires_grad, model.parameters())
|
|
|
|
# the SM model originally follows SGD but Adadelta is used here
|
|
optimizer = torch.optim.Adadelta(parameter, lr=args.lr, weight_decay=args.weight_decay)
|
|
criterion = nn.CrossEntropyLoss()
|
|
early_stop = False
|
|
best_dev_map = 0
|
|
iterations = 0
|
|
iters_not_improved = 0
|
|
epoch = 0
|
|
start = time.time()
|
|
header = ' Time Epoch Iteration Progress (%Epoch) Loss Dev/Loss Accuracy Dev/Accuracy'
|
|
dev_log_template = ' '.join('{:>6.0f},{:>5.0f},{:>9.0f},{:>5.0f}/{:<5.0f} {:>7.0f}%,{:>8.6f},{:8.6f},{:12.4f},{:12.4f}'.split(','))
|
|
log_template = ' '.join('{:>6.0f},{:>5.0f},{:>9.0f},{:>5.0f}/{:<5.0f} {:>7.0f}%,{:>8.6f},{},{:12.4f},{}'.split(','))
|
|
os.makedirs(args.save_path, exist_ok=True)
|
|
os.makedirs(os.path.join(args.save_path, args.dataset), exist_ok=True)
|
|
print(header)
|
|
|
|
index2label = np.array(LABEL.vocab.itos)
|
|
index2qid = np.array(QID.vocab.itos)
|
|
index2question = np.array(ANSWER.vocab.itos)
|
|
|
|
while True:
|
|
if early_stop:
|
|
print("Early Stopping. Epoch: {}, Best Dev Acc: {}".format(epoch, best_dev_map))
|
|
break
|
|
epoch += 1
|
|
train_iter.init_epoch()
|
|
n_correct, n_total = 0, 0
|
|
|
|
for batch_idx, batch in enumerate(train_iter):
|
|
iterations += 1
|
|
model.train(); optimizer.zero_grad()
|
|
scores = model(batch.question, batch.answer, batch.ext_feat)
|
|
n_correct += (torch.max(scores, 1)[1].view(batch.label.size()).data == batch.label.data).sum()
|
|
n_total += batch.batch_size
|
|
train_acc = 100. * n_correct / n_total
|
|
|
|
loss = criterion(scores, batch.label)
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
# Evaluate performance on validation set
|
|
if iterations % args.dev_every == 1:
|
|
# switch model into evaluation mode
|
|
model.eval()
|
|
dev_iter.init_epoch()
|
|
n_dev_correct = 0
|
|
dev_losses = []
|
|
|
|
qids = []
|
|
predictions = []
|
|
labels = []
|
|
for dev_batch_idx, dev_batch in enumerate(dev_iter):
|
|
qid_array = index2qid[np.transpose(dev_batch.qid.cpu().data.numpy())]
|
|
true_label_array = index2label[np.transpose(dev_batch.label.cpu().data.numpy())]
|
|
|
|
scores = model(dev_batch.question, dev_batch.answer, dev_batch.ext_feat)
|
|
n_dev_correct += (torch.max(scores, 1)[1].view(dev_batch.label.size()).data == dev_batch.label.data).sum()
|
|
dev_loss = criterion(scores, dev_batch.label)
|
|
dev_losses.append(dev_loss.data[0])
|
|
index_label = np.transpose(torch.max(scores, 1)[1].view(dev_batch.label.size()).cpu().data.numpy())
|
|
label_array = index2label[index_label]
|
|
# get the relevance scores
|
|
score_array = scores[:, 2].cpu().data.numpy()
|
|
|
|
qids.extend(qid_array.tolist())
|
|
predictions.extend(score_array.tolist())
|
|
labels.extend(true_label_array.tolist())
|
|
|
|
dev_map, dev_mrr = get_map_mrr(qids, predictions, labels)
|
|
print(dev_log_template.format(time.time() - start,
|
|
epoch, iterations, 1 + batch_idx, len(train_iter),
|
|
100. * (1 + batch_idx) / len(train_iter), loss.data[0],
|
|
sum(dev_losses) / len(dev_losses), train_acc, dev_map))
|
|
|
|
# Update validation results
|
|
if dev_map > best_dev_map:
|
|
iters_not_improved = 0
|
|
best_dev_map = dev_map
|
|
snapshot_path = os.path.join(args.save_path, args.dataset, args.mode+'_best_model.pt')
|
|
torch.save(model, snapshot_path)
|
|
else:
|
|
iters_not_improved += 1
|
|
if iters_not_improved >= args.patience:
|
|
early_stop = True
|
|
break
|
|
|
|
if iterations % args.log_every == 1:
|
|
# print progress message
|
|
print(log_template.format(time.time() - start,
|
|
epoch, iterations, 1 + batch_idx, len(train_iter),
|
|
100. * (1 + batch_idx) / len(train_iter), loss.data[0], ' ' * 8,
|
|
n_correct / n_total * 100, ' ' * 12))
|