"""
Copyright (c) 2019-present NAVER Corp.
MIT License
"""
import os
import time
import math
import torch
from model import LMConfig, LanguageModel
class TrainLogger(object):
def __init__(self):
self.init()
def init(self):
self.start = time.time()
self.cnt_add = 0
self.tot_loss = 0.
self.cnt_query = 0
self.cnt_token = 0
def add(self, loss, n_query, n_token):
self.cnt_add += 1
self.tot_loss += loss * n_query
self.cnt_query += n_query
self.cnt_token += n_token
def average(self):
loss_query = self.tot_loss / self.cnt_query if self.cnt_query != 0 else 0.
loss_token = self.tot_loss / self.cnt_token if self.cnt_token != 0 else 0.
return loss_query, loss_token
def elapsed_time(self):
return time.time() - self.start
def print_str(self, time_avg_=False):
loss_query, loss_token = self.average()
time_str = f"{self.elapsed_time() * 1000. / self.cnt_add:6.2f} ms/batch" if time_avg_ else \
f"{self.elapsed_time():6.2f} s"
return f"{time_str} | loss_query {loss_query:6.2f} | token_ppl {math.exp(loss_token):6.2f}"
def get_params(model):
return list(filter(lambda p: p.requires_grad, model.parameters()))
def get_model(model):
return model.module if hasattr(model, 'module') else model
def model_save(path, model, optimizer=None):
model_to_save = get_model(model)
open(os.path.join(path, 'config.json'), 'w').write(str(model_to_save.config))
torch.save(model_to_save.state_dict(), open(os.path.join(path, 'model.pt'), 'wb'))
if optimizer:
torch.save(optimizer.state_dict(), open(os.path.join(path, 'optimizer.pt'), 'wb'))
def model_load(path, model=None, optimizer=None):
config = LMConfig(os.path.join(path, 'config.json'))
if model is None:
model_to_load = LanguageModel(config)
else:
model_to_load = get_model(model)
model_to_load.__init__(config)
model_state_dict = torch.load(open(os.path.join(path, 'model.pt'), 'rb'), map_location=lambda s, l: s)
model_to_load.load_state_dict(model_state_dict)
if optimizer:
optimizer_state_dict = torch.load(open(os.path.join(path, 'optimizer.pt'), 'rb'), map_location=lambda s, l: s)
optimizer.load_state_dict(optimizer_state_dict)
return model_to_load
