| import pdb |
| import os |
| import logging |
| import torch |
| import torch.optim as optim |
| import torch.nn as nn |
| from sklearn import metrics |
|
|
|
|
| class Trainer(): |
| def __init__(self, model, data, params): |
| self.model = model |
| self.data = data |
| self.optimizer = None |
| self.params = params |
|
|
| if params.optimizer == "SGD": |
| self.optimizer = optim.SGD(self.model.parameters(), lr=params.lr, momentum=params.momentum) |
| if params.optimizer == "Adam": |
| self.optimizer = optim.Adam(self.model.parameters(), lr=params.lr) |
|
|
| self.criterion = nn.MarginRankingLoss(self.params.margin, reduction='sum') |
|
|
| self.best_metric = 1e10 |
| self.last_metric = 1e10 |
| self.bad_count = 0 |
|
|
| assert self.optimizer is not None |
|
|
| def one_epoch(self): |
| all_pos_scores = [] |
| all_neg_scores = [] |
| total_loss = 0 |
| for b in range(self.params.nBatches): |
| batch_h, batch_t, batch_r, batch_y = self.data.get_batch(b) |
| batch_h = torch.tensor(batch_h).to(self.params.device) |
| batch_t = torch.tensor(batch_t).to(self.params.device) |
| batch_r = torch.tensor(batch_r).to(self.params.device) |
| batch_y = torch.tensor(batch_y).to(self.params.device) |
| loss, pos_score, neg_score = self.model(batch_h, batch_t, batch_r, batch_y) |
|
|
| all_pos_scores += pos_score.detach().cpu().tolist() |
| all_neg_scores += neg_score.detach().cpu().tolist() |
|
|
| total_loss += loss.detach().cpu() |
|
|
| self.optimizer.zero_grad() |
| loss.backward() |
| self.optimizer.step() |
|
|
| all_labels = [0] * len(all_pos_scores) + [1] * len(all_neg_scores) |
| auc = metrics.roc_auc_score(all_labels, all_pos_scores + all_neg_scores) |
|
|
| return total_loss, auc |
|
|
| def select_model(self, log_data): |
| if log_data['auc'] < self.best_metric: |
| self.bad_count = 0 |
| torch.save(self.model, os.path.join(self.params.exp_dir, 'best_model.pth')) |
| logging.info('Better model found w.r.t MR. Saved it!') |
| self.best_mr = log_data['auc'] |
| else: |
| self.bad_count = self.bad_count + 1 |
| if self.bad_count > self.params.patience: |
| logging.info('Out of patience. Stopping the training loop.') |
| return False |
| self.last_metric = log_data['auc'] |
| return True |
|
|
