File size: 2,481 Bytes
f9c98cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | 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')) # Does it overwrite or fuck with the existing file?
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
|