mirror of
https://github.com/wassname/Castor.git
synced 2026-09-25 13:10:11 +08:00
192 lines
8.2 KiB
Python
192 lines
8.2 KiB
Python
import math
|
|
import time
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch.autograd import Variable
|
|
import torch.optim as optim
|
|
from torch.optim.lr_scheduler import ReduceLROnPlateau
|
|
from scipy.stats import pearsonr, spearmanr
|
|
|
|
# logging setup
|
|
import logging
|
|
logger = logging.getLogger(__name__)
|
|
logger.setLevel(logging.INFO)
|
|
|
|
ch = logging.StreamHandler()
|
|
ch.setLevel(logging.DEBUG)
|
|
formatter = logging.Formatter('%(levelname)s - %(message)s')
|
|
ch.setFormatter(formatter)
|
|
logger.addHandler(ch)
|
|
|
|
|
|
class MPCNNTrainerFactory(object):
|
|
"""
|
|
Get the corresponding Trainer class for a particular dataset.
|
|
"""
|
|
@staticmethod
|
|
def get_trainer(dataset_name, model, optimizer, train_loader, batch_size, sample, log_interval, model_outfile, train_evaluator, test_evaluator, dev_evaluator=None):
|
|
if dataset_name == 'sick':
|
|
return SICKTrainer(model, optimizer, train_loader, batch_size, sample, log_interval, model_outfile, train_evaluator, test_evaluator, dev_evaluator)
|
|
elif dataset_name == 'msrvid':
|
|
return MSRVIDTrainer(model, optimizer, train_loader, batch_size, sample, log_interval, model_outfile, train_evaluator, test_evaluator, dev_evaluator)
|
|
else:
|
|
raise ValueError('{} is not a valid dataset.'.format(dataset_name))
|
|
|
|
|
|
class Trainer(object):
|
|
|
|
"""
|
|
Abstraction for training a model on a Dataset.
|
|
"""
|
|
|
|
def __init__(self, model, optimizer, train_loader, batch_size, sample, log_interval, model_outfile, train_evaluator, test_evaluator, dev_evaluator=None):
|
|
self.model = model
|
|
self.optimizer = optimizer
|
|
self.train_loader = train_loader
|
|
self.batch_size = batch_size
|
|
self.sample = sample
|
|
self.log_interval = log_interval
|
|
self.model_outfile = model_outfile
|
|
self.train_evaluator = train_evaluator
|
|
self.test_evaluator = test_evaluator
|
|
self.dev_evaluator = dev_evaluator
|
|
|
|
def evaluate(self, evaluator, dataset_name):
|
|
scores, metric_names = evaluator.get_scores()
|
|
logger.info('Evaluation metrics for {}:'.format(dataset_name))
|
|
logger.info('\t'.join([' '] + metric_names))
|
|
logger.info('\t'.join([dataset_name] + list(map(str, scores))))
|
|
return scores
|
|
|
|
def train_epoch(self, epoch):
|
|
raise NotImplementedError()
|
|
|
|
def train(self, epochs):
|
|
raise NotImplementedError()
|
|
|
|
|
|
class SICKTrainer(Trainer):
|
|
|
|
def __init__(self, model, optimizer, train_loader, batch_size, sample, log_interval, model_outfile, train_evaluator, test_evaluator, dev_evaluator=None):
|
|
super(SICKTrainer, self).__init__(model, optimizer, train_loader, batch_size, sample, log_interval, model_outfile, train_evaluator, test_evaluator, dev_evaluator)
|
|
|
|
def train_epoch(self, epoch):
|
|
self.model.train()
|
|
total_loss = 0
|
|
for batch_idx, (sentences, labels) in enumerate(self.train_loader):
|
|
sent_a, sent_b = Variable(sentences['a']), Variable(sentences['b'])
|
|
labels = Variable(labels)
|
|
self.optimizer.zero_grad()
|
|
output = self.model(sent_a, sent_b)
|
|
loss = F.kl_div(output, labels)
|
|
total_loss += loss.data[0]
|
|
loss.backward()
|
|
self.optimizer.step()
|
|
if batch_idx % self.log_interval == 0:
|
|
logger.info('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
|
|
epoch, min(batch_idx * self.batch_size, len(self.train_loader.dataset)),
|
|
len(self.train_loader.dataset) if not self.sample else self.sample,
|
|
100. * batch_idx / (len(self.train_loader) if not self.sample else math.ceil(self.sample / self.batch_size)), loss.data[0])
|
|
)
|
|
return total_loss
|
|
|
|
def train(self, epochs):
|
|
scheduler = ReduceLROnPlateau(self.optimizer, mode='max', factor=0.3, patience=2)
|
|
epoch_times = []
|
|
best_dev_score = -1
|
|
for epoch in range(1, epochs + 1):
|
|
start = time.time()
|
|
logger.info('Epoch {} started...'.format(epoch))
|
|
self.train_epoch(epoch)
|
|
|
|
dev_scores = self.evaluate(self.dev_evaluator, 'dev')
|
|
end = time.time()
|
|
duration = end - start
|
|
logger.info('Epoch {} finished in {:.2f} minutes'.format(epoch, duration / 60))
|
|
epoch_times.append(duration)
|
|
|
|
if dev_scores[0] > best_dev_score:
|
|
best_dev_score = dev_scores[0]
|
|
torch.save(self.model, self.model_outfile)
|
|
scheduler.step(dev_scores[0])
|
|
|
|
logger.info('Training took {:.2f} minutes overall...'.format(sum(epoch_times) / 60))
|
|
|
|
|
|
class MSRVIDTrainer(Trainer):
|
|
|
|
def __init__(self, model, optimizer, train_loader, batch_size, sample, log_interval, model_outfile, train_evaluator, test_evaluator, dev_evaluator=None):
|
|
super(MSRVIDTrainer, self).__init__(model, optimizer, train_loader, batch_size, sample, log_interval, model_outfile, train_evaluator, test_evaluator, dev_evaluator)
|
|
|
|
def train_epoch(self, epoch):
|
|
self.model.train()
|
|
|
|
# since MSRVID doesn't have validation set, we manually leave-out some training data for validation
|
|
batches = math.ceil(len(self.train_loader.dataset) / self.batch_size)
|
|
start_val_batch = math.floor(0.8 * batches)
|
|
left_out_val_a, left_out_val_b = [], []
|
|
left_out_val_labels = []
|
|
|
|
for batch_idx, (sentences, labels) in enumerate(self.train_loader):
|
|
sent_a, sent_b = Variable(sentences['a']), Variable(sentences['b'])
|
|
labels = Variable(labels)
|
|
if batch_idx >= start_val_batch:
|
|
left_out_val_a.append(sent_a)
|
|
left_out_val_b.append(sent_b)
|
|
left_out_val_labels.append(labels)
|
|
continue
|
|
self.optimizer.zero_grad()
|
|
output = self.model(sent_a, sent_b)
|
|
loss = F.kl_div(output, labels)
|
|
loss.backward()
|
|
self.optimizer.step()
|
|
if batch_idx % self.log_interval == 0:
|
|
logger.info('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
|
|
epoch, min(batch_idx * self.batch_size, len(self.train_loader.dataset)),
|
|
len(self.train_loader.dataset) if not self.sample else self.sample,
|
|
100. * batch_idx / (len(self.train_loader) if not self.sample else math.ceil(self.sample / self.batch_size)), loss.data[0])
|
|
)
|
|
|
|
self.evaluate(self.train_evaluator, 'train')
|
|
return left_out_val_a, left_out_val_b, left_out_val_labels
|
|
|
|
def train(self, epochs):
|
|
scheduler = ReduceLROnPlateau(self.optimizer, mode='max', factor=0.3, patience=2)
|
|
epoch_times = []
|
|
best_dev_score = -1
|
|
for epoch in range(1, epochs + 1):
|
|
start = time.time()
|
|
logger.info('Epoch {} started...'.format(epoch))
|
|
left_out_a, left_out_b, left_out_label = self.train_epoch(epoch)
|
|
|
|
# manually evaluating the validating set
|
|
left_out_a = torch.cat(left_out_a)
|
|
left_out_b = torch.cat(left_out_b)
|
|
left_out_label = torch.cat(left_out_label)
|
|
output = self.model(left_out_a, left_out_b)
|
|
predict_classes = torch.arange(0, 6).expand(len(left_out_a), 6).cuda()
|
|
true_labels = (predict_classes * left_out_label.data).sum(dim=1)
|
|
predictions = (predict_classes * output.data.exp()).sum(dim=1)
|
|
predictions = predictions.cpu().numpy()
|
|
true_labels = true_labels.cpu().numpy()
|
|
pearson_r = pearsonr(predictions, true_labels)[0]
|
|
for param_group in self.optimizer.param_groups:
|
|
logger.info('Validation size: %s Pearson\'s r: %s', output.size()[0], pearson_r)
|
|
logger.info('Learning rate: %s', param_group['lr'])
|
|
break
|
|
scheduler.step(pearson_r)
|
|
|
|
end = time.time()
|
|
duration = end - start
|
|
logger.info('Epoch {} finished in {:.2f} minutes'.format(epoch, duration / 60))
|
|
epoch_times.append(duration)
|
|
|
|
if pearson_r > best_dev_score:
|
|
best_dev_score = pearson_r
|
|
torch.save(self.model, self.model_outfile)
|
|
|
|
self.evaluate(self.test_evaluator, 'test')
|
|
|
|
logger.info('Training took {:.2f} minutes overall...'.format(sum(epoch_times) / 60))
|