-
Notifications
You must be signed in to change notification settings - Fork 0
/
language_model.py
186 lines (157 loc) · 8.04 KB
/
language_model.py
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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
import math
import pytorch_lightning as pl
import torch
import torch.nn.functional as F
from abc import ABC
from dataclasses import dataclass
from functools import partial
from omegaconf import DictConfig
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from torch import nn, Tensor
from torch.optim.lr_scheduler import LambdaLR, ReduceLROnPlateau
from typing import *
from .padded_tensor import PaddedTensor
from ..train_callbacks import UnconditionalSampler
@dataclass
class LanguageModelHparams(ABC):
batch_size: int = 0 # Just for compatibility with pl.Trainer's auto_scale_batch_size feature
grad_clip_threshold: float = 5.0
init_scale: float = 0.02 # Std. deviation of Gaussian used to initialize the weights
lr: float = 2.5e-4
lr_decay_steps: Optional[int] = 150_000
lr_plateau_patience: Optional[int] = None # If non-None, ReduceLROnPlateau scheduler is used w/ specified patience
warmup_steps: int = 1000
# If beta2 ie set to 0.0, we just use SGD (with momentum set to beta1)- useful for replicating papers that use it
adam_beta1: float = 0.9
adam_beta2: float = 0.999
adam_eps: float = 1e-6 # We need to use 1e-6 since 1e-8 underflows to 0 on fp16
weight_decay: float = 0.01
vocab_size: int = 30522
start_token: Optional[int] = None # If None, it's read off the datamodule's Tokenizer object
end_token: Optional[int] = None
# Whether to divide the loss by the number of tokens in each sequence. This is pretty much always done
# for vanilla language models, but often not done for text VAEs.
divide_loss_by_length: bool = True
early_stopping_metric: str = 'val_nll'
log_samples: bool = True
# noinspection PyMethodMayBeStatic
class LanguageModel(pl.LightningModule, ABC):
def __init__(self, hparams: DictConfig):
super(LanguageModel, self).__init__()
self.save_hyperparameters(hparams)
self.example_input_array = {'batch': {
'token_ids': PaddedTensor.from_raw(torch.zeros(1, 256, dtype=torch.long))}
}
self.tokenizer = None
self.start_token = hparams.start_token
self.end_token = hparams.end_token
def configure_callbacks(self):
metric = self.hparams.early_stopping_metric
callbacks = [
EarlyStopping(monitor=metric, mode='min'),
ModelCheckpoint(monitor=metric, mode='min')
]
return callbacks + [UnconditionalSampler()] if self.hparams.log_samples else callbacks
# The callbacks need to have access to a tokenizer, so let's get a reference to the datamodule's tokenizer.
def setup(self, stage: str):
datamodule = self.trainer.datamodule if stage == 'fit' else self.datamodule
self.tokenizer = datamodule.tokenizer
if not self.start_token and not self.end_token:
vocab = self.tokenizer.get_vocab()
self.start_token = vocab['[CLS]']
self.end_token = vocab['[SEP]']
def configure_optimizers(self, lr: float = None, params = None):
beta1, beta2 = self.hparams.adam_beta1, self.hparams.adam_beta2
if beta2 == 0.0:
opt = torch.optim.SGD(
params or self.parameters(),
lr=lr or self.hparams.lr,
momentum=beta1
)
else:
try:
from deepspeed.ops.adam import FusedAdam as Adam
except ImportError:
print("Couldn't import fused Adam kernel from DeepSpeed, falling back on PyTorch version.")
from torch.optim import Adam
opt = Adam(
params or self.parameters(),
lr=lr or self.hparams.lr,
betas=(beta1, beta2),
weight_decay=self.hparams.weight_decay,
eps=self.hparams.adam_eps
)
# This is mainly just here to reproduce the Lagging Inference Networks paper (He et al. 2019) which uses a
# ReduceLROnPlateau-type learning rate schedule
on_plateau_patience = self.hparams.lr_plateau_patience
if on_plateau_patience is not None:
lr_dict = {
'scheduler': ReduceLROnPlateau(opt, factor=0.5, patience=on_plateau_patience),
'monitor': self.hparams.early_stopping_metric,
'interval': 'epoch'
}
else:
lr_lambda = get_cosine_decay_with_warmup_schedule(self.hparams.lr_decay_steps, self.hparams.warmup_steps)
lr_dict = {
'scheduler': LambdaLR(opt, lr_lambda),
'interval': 'step'
}
return [opt], [lr_dict]
def initialize_weights(self):
# Default BERT weight initialization
for module in self.modules():
if isinstance(module, nn.LayerNorm):
continue
if hasattr(module, 'weight'):
module.weight.data.normal_(0.0, self.hparams.init_scale)
bias = getattr(module, 'bias', None)
bias = getattr(bias, 'data', None)
if bias is not None:
bias.zero_()
# If reduce_batch == False, then this method will not reduce across any dimensions other than sequence length
def stats_from_logits(self, logits: Tensor, labels: Tensor, word_counts: Tensor = None, reduce_batch: bool = True):
nll = F.cross_entropy(input=logits.flatten(end_dim=-2), target=labels.flatten(), ignore_index=0, reduction='none')
nll_sum = nll.view(*logits.shape[:-1]).sum(dim=-1) # Add batch dim(s) back; sum across sequence length
# Divide by the number of non-padding tokens
nll = nll_sum / labels.ne(0).sum(dim=-1) if self.hparams.divide_loss_by_length else nll_sum
if reduce_batch:
# Perplexity is normalized using the number of words, not the number of tokens, so that it's comparable
# across different subword vocabularies
if word_counts is not None:
per_word_nll = (nll_sum / word_counts).mean()
ppl = 2 ** (per_word_nll / math.log(2))
return nll.mean(), ppl.mean()
else:
return nll.mean()
else:
return nll
# These implementations are used by LSTMLanguageModel and TransformerLanguageModel, but are overriden by others
def training_step(self, batch: Dict[str, Tensor], batch_index: int, val: bool = False) -> Tensor:
logits = self.forward(batch) # Remove initial [CLS] token from the labels; autoregressive by default
loss, ppl = self.stats_from_logits(logits, batch['token_ids'][..., 1:], word_counts = batch['num_words'])
self.log('train_nll' if not val else 'val_nll', loss, on_step=not val, on_epoch=val)
self.log('train_ppl' if not val else 'val_ppl', ppl, on_step=not val, on_epoch=val)
return loss
def on_after_backward(self):
grad_norm = nn.utils.clip_grad_norm_(self.parameters(), self.hparams.grad_clip_threshold)
self.log('grad_norm', grad_norm, on_step=True)
def validation_step(self, batch: Dict[str, Tensor], batch_index: int) -> Tensor:
return self.training_step(batch, batch_index, val=True)
# Called by UnconditionalSampler callback
def sample(self, max_length: int, batch_size: int = 1, **kwargs):
return None
# The actual lambda to pass into LambdaLR
def cosine_decay_with_warmup(decay_steps: int, warmup_steps: int, cur_step: int):
if cur_step < warmup_steps: # Warmup phase
return cur_step / warmup_steps
elif not decay_steps: # Just fall back to a constant schedule if we don't know the decay steps
return 1.0
else: # Cosine decay
progress = (cur_step - warmup_steps) / max(1, decay_steps - warmup_steps)
if progress >= 1.0:
print("Learning rate decayed to 0.0. Halting training.")
raise KeyboardInterrupt
return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress)))
# Convenience function
def get_cosine_decay_with_warmup_schedule(decay_steps: int, warmup_steps: int):
return partial(cosine_decay_with_warmup, decay_steps, warmup_steps)