mirror of
https://github.com/wassname/openai-transformer-lm-gutenberg-erotic.git
synced 2026-06-27 16:10:19 +08:00
wassname
27 lines
1012 B
Python
27 lines
1012 B
Python
import torch
|
|
|
|
class LanguageModelingLossCompute:
|
|
" A Loss compute and train function for language modeling tasks."
|
|
def __init__(self, lm_criterion, opt=None):
|
|
self.lm_criterion = lm_criterion
|
|
self.opt = opt
|
|
|
|
# def __call__(self, X, Y, M, clf_logits, lm_logits=None, only_return_losses=False):
|
|
def __call__(self, X, Y, M, lm_logits, only_return_losses=False):
|
|
# Language modeling loss
|
|
x_shifted = X[:, 1:, 0].contiguous().view(-1)
|
|
M = M.view(-1, M.size(-1))
|
|
lm_losses = self.lm_criterion(lm_logits, x_shifted)
|
|
lm_losses = lm_losses.view(X.size(0), X.size(-2) - 1)
|
|
lm_losses = lm_losses * M[:, 1:]
|
|
lm_losses = lm_losses.sum(1) / torch.sum(M[:, 1:], 1)
|
|
if only_return_losses:
|
|
return lm_losses
|
|
|
|
train_loss = lm_losses.sum()
|
|
train_loss.backward()
|
|
if self.opt is not None:
|
|
self.opt.step()
|
|
self.opt.zero_grad()
|
|
return train_loss.item()
|