chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,90 @@
|
||||
# Adaptive Span
|
||||
|
||||
Adaptive Span is a novel self-attention mechanism that can learn its optimal
|
||||
attention span. This allows us to extend significantly the maximum context size
|
||||
used in Transformer, while maintaining control over their memory footprint
|
||||
and computational time. It uses the Truncated BPTT technique for training,
|
||||
as in [transformerXL](https://github.com/pytorch/fairseq/blob/master/examples/truncated_bptt/README.md).
|
||||
|
||||
Adaptive Span was introduced by paper:
|
||||
[Adaptive Attention Span in Transformers](https://arxiv.org/abs/1905.07799),
|
||||
which achieved state-of-the-art language modeling results at the time of publication.
|
||||
|
||||
We manage to reproduce their result in fairseq and keep most of the
|
||||
[original implementation](https://github.com/facebookresearch/adaptive-span) untouched.
|
||||
You can refer to the their sweep file as well if any combination of hyperparameter is not clear.
|
||||
|
||||
##### 0. Setup
|
||||
|
||||
First you need to process the Enwik8 dataset, we use the pre-tokenized dataset
|
||||
from [adaptive span paper](https://github.com/facebookresearch/adaptive-span/blob/master/get_data.sh).
|
||||
You can download the dataset, and then run:
|
||||
```bash
|
||||
fairseq-preprocess --only-source --trainpref ~/data/enwik8/train.txt \
|
||||
--validpref ~/data/enwik8/valid.txt --testpref ~/data/enwik8/test.txt \
|
||||
--destdir ~/data/enwik8/data-bin/ --joined-dictionary --workers 20
|
||||
```
|
||||
|
||||
##### 1. Train a Adaptive Span model on Enwik8
|
||||
|
||||
We will train a 12-layer Adaptive Span model following the [hyperparameters
|
||||
used in the original
|
||||
paper](https://github.com/facebookresearch/adaptive-span/blob/master/experiments/enwik8.sh).
|
||||
|
||||
The following command assumes 4 GPUs, so that the total batch size is 64
|
||||
sequences (4 x 16). Training should take 2-3 days on 4 V100 GPUs:
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 fairseq-train \
|
||||
--user-dir examples/adaptive_span \
|
||||
--data ~/data/enwik8/data-bin/ \
|
||||
--fp16 --fp16-no-flatten-grads --max-update 600000 \
|
||||
--task truncated_bptt_lm --tokens-per-sample 512 --arch adaptive_span \
|
||||
--n-layer 12 --d-model 512 --n-head 8 --d-inner 2048 --dropout 0.3 \
|
||||
--attn-span 8192 --optimizer adagrad_with_grad_clip --adagrad-clip 0.03 \
|
||||
--validate-interval-updates 1000 \
|
||||
--lr-scheduler fixed --warmup-updates 32000 --batch-size-valid 32 \
|
||||
--lr 0.07 --criterion adaptive_span_loss --batch-size 16 --update-freq 1 \
|
||||
--seed 2 --log-format json --log-interval 25 --aux-loss-scaler 5e-07
|
||||
```
|
||||
This should land around 1.05 on validation, 1.03 on test. You can lower the
|
||||
--aux-loss-scaler for better performance (longer span). It gives ~0.03 bpc
|
||||
improvement to the transformerXL baseline here.
|
||||
If training on a single GPU, set `--update-freq=4` to accumulate 4x gradients
|
||||
and simulate training on 4 GPUs.
|
||||
You can also reproduce the transformerXL result on enwik8 using this code base.
|
||||
It should land around 1.06 on test,matching the [original paper](https://github.com/kimiyoung/transformer-xl/blob/master/pytorch/run_enwik8_base.sh).
|
||||
You can try by
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 fairseq-train \
|
||||
--user-dir examples/truncated_bptt \
|
||||
~/data/enwik8/data-bin/ \
|
||||
--task truncated_bptt_lm --fp16 --max-update 400000 \
|
||||
--tokens-per-sample 512 --arch transformer_xl --n-layer 12 \
|
||||
--d-model 512 --n-head 8 --d-head 64 --d-inner 2048 --dropout 0.1 \
|
||||
--dropatt 0.0 --mem-len 512 --optimizer adam --clip-norm 0.25 \
|
||||
--lr-scheduler cosine --warmup-updates 0 \
|
||||
--lr 0.0 --lr 0.00025 --batch-size 15 \
|
||||
--update-freq 1 --seed 2 --log-format json --log-interval 25 \
|
||||
--fp16
|
||||
```
|
||||
|
||||
##### 2. Evaluate
|
||||
For Adaptive Span:
|
||||
```bash
|
||||
fairseq-eval-lm ~/data/enwik8/data-bin/ --path model/checkpoint_best.pt \
|
||||
--user-dir examples/adaptive_span \
|
||||
--task truncated_bptt_lm --batch-size 8 --tokens-per-sample 512 --gen-subset test
|
||||
```
|
||||
For Transformer-XL evaluation:
|
||||
```bash
|
||||
fairseq-eval-lm ~/data/enwik8/data-bin/ --path model/checkpoint_best.pt \
|
||||
--user-dir examples/truncated_bptt/ --task truncated_bptt_lm --batch-size 8 \
|
||||
--tokens-per-sample 80 \
|
||||
--model-overrides '{"mem_len":2100,"clamp_len":820,"same_length":True}' \
|
||||
--gen-subset valid
|
||||
```
|
||||
|
||||
*Note:* During training the model saw 512 tokens of context
|
||||
(``--tokens-per-sample=512``), with batch size 8. These settings match the evaluation
|
||||
settings from [the original
|
||||
paper](https://github.com/facebookresearch/adaptive-span/blob/master/experiments/enwik8.sh).
|
||||
@@ -0,0 +1,19 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import importlib
|
||||
import os
|
||||
|
||||
# automatically import any Python files in the current directory
|
||||
cur_dir = os.path.dirname(__file__)
|
||||
for file in os.listdir(cur_dir):
|
||||
path = os.path.join(cur_dir, file)
|
||||
if (
|
||||
not file.startswith("_")
|
||||
and not file.startswith(".")
|
||||
and (file.endswith(".py") or os.path.isdir(path))
|
||||
):
|
||||
mod_name = file[: file.find(".py")] if file.endswith(".py") else file
|
||||
module = importlib.import_module(__name__ + "." + mod_name)
|
||||
@@ -0,0 +1,128 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from torch.optim import Adagrad
|
||||
|
||||
from fairseq.optim import LegacyFairseqOptimizer, register_optimizer
|
||||
|
||||
|
||||
@register_optimizer("adagrad_with_grad_clip")
|
||||
class FairseqAdagradWithGradClip(LegacyFairseqOptimizer):
|
||||
def __init__(self, args, params):
|
||||
super().__init__(args)
|
||||
self._optimizer = AdagradWithGradClip(params, **self.optimizer_config)
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
"""Add optimizer-specific arguments to the parser."""
|
||||
# fmt: off
|
||||
parser.add_argument('--weight-decay', '--wd', default=0.0, type=float, metavar='WD',
|
||||
help='weight decay')
|
||||
parser.add_argument('--adagrad-clip', default=0.0, type=float, metavar='D',
|
||||
help='internal grad clip')
|
||||
# fmt: on
|
||||
|
||||
@property
|
||||
def optimizer_config(self):
|
||||
"""
|
||||
Return a kwarg dictionary that will be used to override optimizer
|
||||
args stored in checkpoints. This allows us to load a checkpoint and
|
||||
resume training using a different set of optimizer args, e.g., with a
|
||||
different learning rate.
|
||||
"""
|
||||
return {
|
||||
"lr": self.args.lr[0],
|
||||
"weight_decay": self.args.weight_decay,
|
||||
"grad_clip": self.args.adagrad_clip,
|
||||
}
|
||||
|
||||
@property
|
||||
def supports_flat_params(self):
|
||||
return False
|
||||
|
||||
|
||||
def _clip_grad(clr, grad, group_grad_clip):
|
||||
if group_grad_clip > 0:
|
||||
norm = grad.norm(2).item()
|
||||
if norm > group_grad_clip:
|
||||
clr *= group_grad_clip / (norm + 1e-10)
|
||||
return clr
|
||||
|
||||
|
||||
class AdagradWithGradClip(Adagrad):
|
||||
"""Adagrad algorithm with custom gradient clipping"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
params,
|
||||
lr=1e-2,
|
||||
lr_decay=0,
|
||||
weight_decay=0,
|
||||
initial_accumulator_value=0,
|
||||
grad_clip=0,
|
||||
):
|
||||
Adagrad.__init__(
|
||||
self,
|
||||
params,
|
||||
lr=lr,
|
||||
lr_decay=lr_decay,
|
||||
weight_decay=weight_decay,
|
||||
initial_accumulator_value=initial_accumulator_value,
|
||||
)
|
||||
self.defaults["grad_clip"] = grad_clip
|
||||
self.param_groups[0].setdefault("grad_clip", grad_clip)
|
||||
|
||||
def step(self, closure=None):
|
||||
loss = None
|
||||
if closure is not None:
|
||||
loss = closure()
|
||||
|
||||
for group in self.param_groups:
|
||||
for p in group["params"]:
|
||||
if p.grad is None:
|
||||
continue
|
||||
|
||||
grad = p.grad.data
|
||||
state = self.state[p]
|
||||
|
||||
state["step"] += 1
|
||||
|
||||
if group["weight_decay"] != 0:
|
||||
if p.grad.data.is_sparse:
|
||||
raise RuntimeError(
|
||||
"weight_decay option is "
|
||||
"not compatible with sparse "
|
||||
"gradients"
|
||||
)
|
||||
grad = grad.add(group["weight_decay"], p.data)
|
||||
|
||||
clr = group["lr"] / (1 + (state["step"] - 1) * group["lr_decay"])
|
||||
|
||||
# clip
|
||||
clr = _clip_grad(clr=clr, grad=grad, group_grad_clip=group["grad_clip"])
|
||||
|
||||
if grad.is_sparse:
|
||||
# the update is non-linear so indices must be unique
|
||||
grad = grad.coalesce()
|
||||
grad_indices = grad._indices()
|
||||
grad_values = grad._values()
|
||||
size = grad.size()
|
||||
|
||||
def make_sparse(values):
|
||||
constructor = grad.new
|
||||
if grad_indices.dim() == 0 or values.dim() == 0:
|
||||
return constructor().resize_as_(grad)
|
||||
return constructor(grad_indices, values, size)
|
||||
|
||||
state["sum"].add_(make_sparse(grad_values.pow(2)))
|
||||
std = state["sum"]._sparse_mask(grad)
|
||||
std_values = std._values().sqrt_().add_(1e-10)
|
||||
p.data.add_(-clr, make_sparse(grad_values / std_values))
|
||||
else:
|
||||
state["sum"].addcmul_(1, grad, grad)
|
||||
std = state["sum"].sqrt().add_(1e-10)
|
||||
p.data.addcdiv_(-clr, grad, std)
|
||||
|
||||
return loss
|
||||
@@ -0,0 +1,160 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class AdaptiveMask(nn.Module):
|
||||
"""Soft masking function for adaptive size.
|
||||
It masks out the last K values of an input. The masking value
|
||||
goes from 1 to 0 gradually, so K can be learned with
|
||||
back-propagation.
|
||||
Args:
|
||||
max_size: maximum size (i.e. input dimension)
|
||||
ramp_size: size of the ramp going from 0 to 1
|
||||
init_val: initial size proportion not to be masked out
|
||||
shape: learn multiple sizes independent of each other
|
||||
"""
|
||||
|
||||
def __init__(self, max_size, ramp_size, init_val=0, shape=(1,)):
|
||||
nn.Module.__init__(self)
|
||||
self._max_size = max_size
|
||||
self._ramp_size = ramp_size
|
||||
self.current_val = nn.Parameter(torch.zeros(*shape) + init_val)
|
||||
mask_template = torch.linspace(1 - max_size, 0, steps=max_size)
|
||||
self.register_buffer("mask_template", mask_template)
|
||||
|
||||
def forward(self, x):
|
||||
mask = self.mask_template.float() + self.current_val.float() * self._max_size
|
||||
mask = mask / self._ramp_size + 1
|
||||
mask = mask.clamp(0, 1)
|
||||
if x.size(-1) < self._max_size:
|
||||
# the input could have been trimmed beforehand to save computation
|
||||
mask = mask.narrow(-1, self._max_size - x.size(-1), x.size(-1))
|
||||
x = (x * mask).type_as(x)
|
||||
return x
|
||||
|
||||
def get_current_max_size(self, include_ramp=True):
|
||||
current_size = math.ceil(self.current_val.max().item() * self._max_size)
|
||||
if include_ramp:
|
||||
current_size += self._ramp_size
|
||||
current_size = max(0, min(self._max_size, current_size))
|
||||
return current_size
|
||||
|
||||
def get_current_avg_size(self, include_ramp=True):
|
||||
current_size = math.ceil(
|
||||
self.current_val.float().mean().item() * self._max_size
|
||||
)
|
||||
if include_ramp:
|
||||
current_size += self._ramp_size
|
||||
current_size = max(0, min(self._max_size, current_size))
|
||||
return current_size
|
||||
|
||||
def clamp_param(self):
|
||||
"""this need to be called after each update"""
|
||||
self.current_val.data.clamp_(0, 1)
|
||||
|
||||
|
||||
class AdaptiveSpan(nn.Module):
|
||||
"""Adaptive attention span for Transformerself.
|
||||
This module learns an attention span length from data for each
|
||||
self-attention head.
|
||||
Args:
|
||||
attn_span: maximum attention span
|
||||
adapt_span_loss: loss coefficient for the span length
|
||||
adapt_span_ramp: length of the masking ramp
|
||||
adapt_span_init: initial size ratio
|
||||
adapt_span_cache: adapt cache size to reduce memory usage
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
attn_span,
|
||||
adapt_span_ramp,
|
||||
adapt_span_init,
|
||||
n_head,
|
||||
adapt_span_layer,
|
||||
**kargs
|
||||
):
|
||||
nn.Module.__init__(self)
|
||||
self._max_span = attn_span
|
||||
self._n_head = n_head
|
||||
self._adapt_span_layer = adapt_span_layer
|
||||
if self._adapt_span_layer:
|
||||
self._mask = AdaptiveMask(
|
||||
max_size=self._max_span,
|
||||
ramp_size=adapt_span_ramp,
|
||||
init_val=adapt_span_init,
|
||||
)
|
||||
else:
|
||||
self._mask = AdaptiveMask(
|
||||
max_size=self._max_span,
|
||||
ramp_size=adapt_span_ramp,
|
||||
init_val=adapt_span_init,
|
||||
shape=(n_head, 1, 1),
|
||||
)
|
||||
|
||||
def forward(self, attn, normalize=True):
|
||||
"""mask attention with the right span"""
|
||||
# batch and head dimensions are merged together, so separate them first
|
||||
self.clamp_param()
|
||||
if self._adapt_span_layer:
|
||||
attn = self._mask(attn)
|
||||
else:
|
||||
B = attn.size(0) # batch size
|
||||
M = attn.size(1) # block size
|
||||
attn = attn.reshape(B // self._n_head, self._n_head, M, -1)
|
||||
attn = self._mask(attn)
|
||||
attn = attn.view(B, M, -1)
|
||||
return attn
|
||||
|
||||
def get_trim_len(self):
|
||||
"""how much of memory can be trimmed to reduce computation"""
|
||||
L = self._max_span
|
||||
trim_len = min(L - 1, L - self._mask.get_current_max_size())
|
||||
# too fine granularity might be bad for the memory management
|
||||
trim_len = math.floor(trim_len / 64) * 64
|
||||
return trim_len
|
||||
|
||||
def trim_memory(self, query, key, value, key_pe):
|
||||
"""trim out unnecessary memory beforehand to reduce computation"""
|
||||
trim_len = self.get_trim_len()
|
||||
cache_size = key.size(1) - query.size(1)
|
||||
trim_len_cache = trim_len - (self._max_span - cache_size)
|
||||
if trim_len_cache > 0:
|
||||
key = key[:, trim_len_cache:, :]
|
||||
value = value[:, trim_len_cache:, :]
|
||||
elif trim_len_cache < 0:
|
||||
# cache is too short! this happens when validation resumes
|
||||
# after a lot of updates.
|
||||
key = F.pad(key, [0, 0, -trim_len_cache, 0])
|
||||
value = F.pad(value, [0, 0, -trim_len_cache, 0])
|
||||
if trim_len > 0:
|
||||
if key_pe is not None:
|
||||
key_pe = key_pe[:, :, trim_len:]
|
||||
return key, value, key_pe
|
||||
|
||||
def get_cache_size(self):
|
||||
"""determine how long the cache should be"""
|
||||
trim_len = self.get_trim_len()
|
||||
# give a buffer of 64 steps since a span might increase
|
||||
# in future updates
|
||||
return min(self._max_span, self._max_span - trim_len + 64)
|
||||
|
||||
def get_loss(self):
|
||||
"""a loss term for regularizing the span length"""
|
||||
return self._max_span * self._mask.current_val.float().mean()
|
||||
|
||||
def get_current_max_span(self):
|
||||
return self._mask.get_current_max_size()
|
||||
|
||||
def get_current_avg_span(self):
|
||||
return self._mask.get_current_avg_size()
|
||||
|
||||
def clamp_param(self):
|
||||
self._mask.clamp_param()
|
||||
@@ -0,0 +1,106 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch.nn.functional as F
|
||||
from fairseq import metrics, utils
|
||||
from fairseq.criterions import register_criterion
|
||||
from fairseq.criterions.cross_entropy import CrossEntropyCriterion
|
||||
from fairseq.dataclass import FairseqDataclass
|
||||
from omegaconf import II
|
||||
|
||||
|
||||
@dataclass
|
||||
class AdaptiveSpanCriterionConfig(FairseqDataclass):
|
||||
sentence_avg: bool = II("optimization.sentence_avg")
|
||||
|
||||
|
||||
@register_criterion("adaptive_span_loss", dataclass=AdaptiveSpanCriterionConfig)
|
||||
class AdaptiveSpanCriterion(CrossEntropyCriterion):
|
||||
def __init__(self, task, sentence_avg):
|
||||
super().__init__(task, sentence_avg)
|
||||
|
||||
def forward(self, model, sample, reduce=True):
|
||||
"""Compute the loss for the given sample.
|
||||
|
||||
Returns a tuple with three elements:
|
||||
1) the loss here is summed, different from the adaptive span code
|
||||
2) the sample size, which is used as the denominator for the gradient
|
||||
3) logging outputs to display while training
|
||||
"""
|
||||
net_output = model(**sample["net_input"])
|
||||
loss, aux_loss, avg_span, max_span = self.compute_loss(
|
||||
model, net_output, sample, reduce=reduce
|
||||
)
|
||||
sample_size = (
|
||||
sample["target"].size(0) if self.sentence_avg else sample["ntokens"]
|
||||
)
|
||||
loss /= sample_size
|
||||
total_loss = loss + aux_loss
|
||||
sample_size = 1
|
||||
|
||||
logging_output = {
|
||||
"loss": loss.data,
|
||||
"ntokens": sample["ntokens"],
|
||||
"nsentences": sample["target"].size(0),
|
||||
"sample_size": sample_size,
|
||||
"total_loss": total_loss.data,
|
||||
"avg_span": avg_span * sample_size,
|
||||
"max_span": max_span * sample_size,
|
||||
}
|
||||
return total_loss, sample_size, logging_output
|
||||
|
||||
def compute_loss(self, model, net_output, sample, reduce=True):
|
||||
loss, _ = super().compute_loss(model, net_output, sample, reduce)
|
||||
aux_loss = model.get_aux_loss()
|
||||
avg_span = model.get_current_avg_span()
|
||||
max_span = model.get_current_max_span()
|
||||
return loss, aux_loss, avg_span, max_span
|
||||
|
||||
@staticmethod
|
||||
def reduce_metrics(logging_outputs) -> None:
|
||||
"""Aggregate logging outputs from data parallel training."""
|
||||
loss_sum = sum(log.get("loss", 0) for log in logging_outputs)
|
||||
ntokens = sum(log.get("ntokens", 0) for log in logging_outputs)
|
||||
sample_size = sum(log.get("sample_size", 0) for log in logging_outputs)
|
||||
total_loss_sum = sum(log.get("total_loss", 0) for log in logging_outputs)
|
||||
avg_span_sum = sum(log.get("avg_span", 0) for log in logging_outputs)
|
||||
max_span_sum = sum(log.get("max_span", 0) for log in logging_outputs)
|
||||
|
||||
# we divide by log(2) to convert the loss from base e to base 2
|
||||
metrics.log_scalar(
|
||||
"loss", loss_sum / sample_size / math.log(2), sample_size, round=3
|
||||
)
|
||||
metrics.log_scalar("avg_span", avg_span_sum / sample_size, sample_size, round=3)
|
||||
metrics.log_scalar("max_span", max_span_sum / sample_size, sample_size, round=3)
|
||||
# total loss contains the L1 norm on adaptive-span
|
||||
metrics.log_scalar(
|
||||
"total_loss",
|
||||
total_loss_sum / sample_size / math.log(2),
|
||||
sample_size,
|
||||
round=3,
|
||||
)
|
||||
if sample_size != ntokens:
|
||||
metrics.log_scalar(
|
||||
"nll_loss", loss_sum / ntokens / math.log(2), ntokens, round=3
|
||||
)
|
||||
metrics.log_derived(
|
||||
"ppl", lambda meters: utils.get_perplexity(meters["nll_loss"].avg)
|
||||
)
|
||||
else:
|
||||
metrics.log_derived(
|
||||
"ppl", lambda meters: utils.get_perplexity(meters["loss"].avg)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def logging_outputs_can_be_summed() -> bool:
|
||||
"""
|
||||
Whether the logging outputs returned by `forward` can be summed
|
||||
across workers prior to calling `reduce_metrics`. Setting this
|
||||
to True will improves distributed training speed.
|
||||
"""
|
||||
return True
|
||||
@@ -0,0 +1,263 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from fairseq.modules.layer_norm import LayerNorm
|
||||
|
||||
from .adaptive_span_attention import AdaptiveSpan
|
||||
|
||||
# Size notations:
|
||||
# B = batch_size, H = d_model, M = block_size, L = attn_span
|
||||
|
||||
|
||||
def _skew(X, pad_value):
|
||||
"""shift every row 1 step to right"""
|
||||
# X = B x M x L
|
||||
B, M, L = X.size()
|
||||
X = F.pad(X, (0, M + 1), value=pad_value) # B x M x (L+M+1)
|
||||
X = X.view(B, -1) # B x ML+MM+M
|
||||
X = X[:, :-M] # B x ML+MM
|
||||
X = X.view(B, M, M + L) # B x M x L+M
|
||||
return X
|
||||
|
||||
|
||||
def _unskew(X):
|
||||
"""reverse _skew operation"""
|
||||
# X = B x M x L+M
|
||||
B, M, L = X.size()
|
||||
L -= M
|
||||
X = X.view(B, -1) # B x ML+MM
|
||||
X = F.pad(X, (0, M)) # B x ML+MM+M
|
||||
X = X.view(B, M, M + L + 1) # B x M x L+M+1
|
||||
X = X[:, :, :L] # B x M x L
|
||||
return X
|
||||
|
||||
|
||||
class SeqAttention(nn.Module):
|
||||
"""Sequential self-attention layer.
|
||||
Each token will attend to its previous fixed number of steps.
|
||||
Note that attention doesn't include the current step itself.
|
||||
"""
|
||||
|
||||
def __init__(self, d_model, n_head, attn_span, dropout, adapt_span_layer, **kargs):
|
||||
nn.Module.__init__(self)
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
self.d_model = d_model # size of a single head
|
||||
self.attn_span = attn_span
|
||||
self.adaptive_span = AdaptiveSpan(
|
||||
attn_span=attn_span,
|
||||
n_head=n_head,
|
||||
adapt_span_layer=adapt_span_layer,
|
||||
**kargs
|
||||
)
|
||||
|
||||
def forward(self, query, key, value, key_pe):
|
||||
# query size = B x M x H
|
||||
# key, value sizes = B x (M+L) x H
|
||||
|
||||
key, value, key_pe = self.adaptive_span.trim_memory(query, key, value, key_pe)
|
||||
|
||||
# compute attention from context
|
||||
# B x M (dest) x (M+L) (src)
|
||||
attn_cont = torch.matmul(query, key.transpose(-1, -2))
|
||||
attn_cont = _unskew(attn_cont) # B x M x L
|
||||
|
||||
# compute the effect of position embedding
|
||||
attn_pos = torch.matmul(query, key_pe) # B x M x L_pos
|
||||
attn = attn_cont + attn_pos
|
||||
|
||||
attn = attn / math.sqrt(self.d_model) # B x M X L_pos
|
||||
|
||||
attn = F.softmax(attn.float(), dim=-1).type_as(attn)
|
||||
|
||||
# trim attention lengths according to the learned span
|
||||
attn = self.adaptive_span(attn)
|
||||
|
||||
attn = self.dropout(attn) # B x M X L_pos
|
||||
|
||||
attn_cont = _skew(attn, 0) # B x M X (L+M)
|
||||
out = torch.matmul(attn_cont, value) # B x M x H
|
||||
return out
|
||||
|
||||
def get_cache_size(self):
|
||||
return self.adaptive_span.get_cache_size()
|
||||
|
||||
|
||||
class MultiHeadSeqAttention(nn.Module):
|
||||
def __init__(self, d_model, n_head, **kargs):
|
||||
nn.Module.__init__(self)
|
||||
assert d_model % n_head == 0
|
||||
self.n_head = n_head
|
||||
self.head_dim = d_model // n_head
|
||||
self.attn = SeqAttention(d_model=self.head_dim, n_head=n_head, **kargs)
|
||||
self.proj_query = nn.Linear(d_model, d_model, bias=False)
|
||||
nn.init.xavier_normal_(self.proj_query.weight)
|
||||
self.proj_out = nn.Linear(d_model, d_model, bias=False)
|
||||
nn.init.xavier_normal_(self.proj_out.weight)
|
||||
self.proj_val = nn.Linear(d_model, d_model, bias=False)
|
||||
nn.init.xavier_normal_(self.proj_val.weight)
|
||||
self.proj_key = nn.Linear(d_model, d_model, bias=False)
|
||||
nn.init.xavier_normal_(self.proj_key.weight)
|
||||
|
||||
def head_reshape(self, x):
|
||||
K = self.n_head
|
||||
D = self.head_dim
|
||||
x = x.view(x.size()[:-1] + (K, D)) # B x (M+L) x K x D
|
||||
x = x.transpose(1, 2).contiguous() # B x K x (M+L) x D
|
||||
x = x.view(-1, x.size(-2), x.size(-1)) # B_K x (M+L) x D
|
||||
return x
|
||||
|
||||
def forward(self, query, key, value, key_pe):
|
||||
B = query.size(0)
|
||||
K = self.n_head
|
||||
D = self.head_dim
|
||||
M = query.size(1)
|
||||
|
||||
query = self.proj_query(query)
|
||||
query = self.head_reshape(query)
|
||||
value = self.proj_val(value)
|
||||
value = self.head_reshape(value)
|
||||
key = self.proj_key(key)
|
||||
key = self.head_reshape(key)
|
||||
|
||||
out = self.attn(query, key, value, key_pe) # B_K x M x D
|
||||
out = out.view(B, K, M, D) # B x K x M x D
|
||||
out = out.transpose(1, 2).contiguous() # B x M x K x D
|
||||
out = out.view(B, M, -1) # B x M x K_D
|
||||
out = self.proj_out(out)
|
||||
return out
|
||||
|
||||
|
||||
class FeedForwardLayer(nn.Module):
|
||||
def __init__(self, d_model, d_inner, dropout, **kargs):
|
||||
nn.Module.__init__(self)
|
||||
self.fc1 = nn.Linear(d_model, d_inner)
|
||||
self.fc2 = nn.Linear(d_inner, d_model)
|
||||
nn.init.xavier_uniform_(self.fc1.weight)
|
||||
nn.init.xavier_uniform_(self.fc2.weight)
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
|
||||
def forward(self, h):
|
||||
h1 = F.relu(self.fc1(h))
|
||||
h1 = self.dropout(h1)
|
||||
h2 = self.fc2(h1)
|
||||
return h2
|
||||
|
||||
|
||||
class TransformerSeqLayer(nn.Module):
|
||||
def __init__(self, d_model, **kargs):
|
||||
nn.Module.__init__(self)
|
||||
self.attn = MultiHeadSeqAttention(d_model=d_model, **kargs)
|
||||
self.norm1 = LayerNorm(d_model)
|
||||
self.ff = FeedForwardLayer(d_model=d_model, **kargs)
|
||||
self.norm2 = LayerNorm(d_model)
|
||||
|
||||
def forward(self, h, h_cache, key_pe):
|
||||
# h = B x M x H
|
||||
# h_cache = B x L x H
|
||||
h_all = torch.cat([h_cache, h], dim=1) # B x (M+L) x H
|
||||
attn_out = self.attn(h, h_all, h_all, key_pe)
|
||||
h = self.norm1(h + attn_out) # B x M x H
|
||||
if self.ff is not None:
|
||||
ff_out = self.ff(h)
|
||||
out = self.norm2(h + ff_out) # B x M x H
|
||||
else:
|
||||
out = h
|
||||
return out
|
||||
|
||||
def get_cache_size(self):
|
||||
return self.attn.attn.get_cache_size()
|
||||
|
||||
|
||||
class TransformerSeq(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
vocab_size,
|
||||
d_model,
|
||||
n_head,
|
||||
n_layer,
|
||||
attn_span,
|
||||
emb_dropout,
|
||||
aux_loss_scaler,
|
||||
adapt_span_layer,
|
||||
**kargs
|
||||
):
|
||||
nn.Module.__init__(self)
|
||||
# token embeddings
|
||||
self.in_emb = nn.Embedding(vocab_size, d_model)
|
||||
nn.init.normal_(self.in_emb.weight, mean=0, std=d_model ** -0.5)
|
||||
self.out_emb = nn.Linear(d_model, vocab_size)
|
||||
self.aux_loss_scaler = aux_loss_scaler
|
||||
if emb_dropout > 0:
|
||||
self.emb_dropout = nn.Dropout(emb_dropout)
|
||||
else:
|
||||
self.emb_dropout = None
|
||||
# position embeddings
|
||||
self.key_pe = nn.Parameter(torch.randn(1, d_model // n_head, attn_span))
|
||||
|
||||
self.layers = nn.ModuleList()
|
||||
self.layers.extend(
|
||||
TransformerSeqLayer(
|
||||
d_model=d_model,
|
||||
n_head=n_head,
|
||||
attn_span=attn_span,
|
||||
adapt_span_layer=adapt_span_layer,
|
||||
**kargs
|
||||
)
|
||||
for _ in range(n_layer)
|
||||
)
|
||||
|
||||
def forward(self, x, h_cache, target=None):
|
||||
# x size = B x M
|
||||
block_size = x.size(1)
|
||||
h = self.in_emb(x) # B x M x H
|
||||
if self.emb_dropout is not None:
|
||||
h = self.emb_dropout(h)
|
||||
|
||||
h_cache_next = []
|
||||
for l, layer in enumerate(self.layers):
|
||||
cache_size = layer.attn.attn.get_cache_size()
|
||||
if cache_size > block_size:
|
||||
h_cache_next_l = torch.cat(
|
||||
[h_cache[l][:, -cache_size + block_size :, :], h], dim=1
|
||||
).detach()
|
||||
else:
|
||||
h_cache_next_l = h[:, -cache_size:, :].detach()
|
||||
h_cache_next.append(h_cache_next_l)
|
||||
h = layer(h, h_cache[l], self.key_pe) # B x M x H
|
||||
|
||||
if self.emb_dropout is not None:
|
||||
h = self.emb_dropout(h)
|
||||
|
||||
out = F.log_softmax(self.out_emb(h).float(), dim=-1).type_as(h)
|
||||
dummy_loss = None
|
||||
|
||||
return out, h_cache_next, dummy_loss
|
||||
|
||||
def get_aux_loss(self):
|
||||
loss = 0.0
|
||||
for layer in self.layers:
|
||||
loss += layer.attn.attn.adaptive_span.get_loss()
|
||||
return self.aux_loss_scaler * loss
|
||||
|
||||
def get_current_max_span(self):
|
||||
max_span = 0.0
|
||||
for layer in self.layers:
|
||||
max_span = max(
|
||||
max_span, layer.attn.attn.adaptive_span.get_current_max_span()
|
||||
)
|
||||
return max_span
|
||||
|
||||
def get_current_avg_span(self):
|
||||
avg_span = 0.0
|
||||
for layer in self.layers:
|
||||
avg_span += layer.attn.attn.adaptive_span.get_current_avg_span()
|
||||
return avg_span / len(self.layers)
|
||||
@@ -0,0 +1,145 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import torch
|
||||
from fairseq.dataclass import FairseqDataclass
|
||||
from fairseq.models import (
|
||||
FairseqIncrementalDecoder,
|
||||
FairseqLanguageModel,
|
||||
register_model,
|
||||
)
|
||||
from .adaptive_span_model import TransformerSeq as AdaptiveSpanTransformerModel
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AdaptiveSpanSmallConfig(FairseqDataclass):
|
||||
# defaults come from https://github.com/facebookresearch/adaptive-span/blob/master/experiments/enwik8_small.sh
|
||||
vocab_size: int = 50
|
||||
d_model: int = 256
|
||||
n_head: int = 4
|
||||
d_inner: int = 1024
|
||||
n_layer: int = 8
|
||||
attn_span: int = 1024
|
||||
dropout: float = 0.0
|
||||
emb_dropout: float = 0.0
|
||||
adapt_span_ramp: int = 32
|
||||
adapt_span_init: float = 0.0
|
||||
aux_loss_scaler: float = 0.000002
|
||||
adapt_span_layer: bool = False
|
||||
|
||||
|
||||
@register_model("adaptive_span", dataclass=AdaptiveSpanSmallConfig)
|
||||
class AdaptiveSpanTransformer(FairseqLanguageModel):
|
||||
@classmethod
|
||||
def build_model(cls, cfg: AdaptiveSpanSmallConfig, task):
|
||||
return cls(AdaptiveSpanDecoder(cfg, task))
|
||||
|
||||
def get_aux_loss(self):
|
||||
return self.decoder.get_aux_loss()
|
||||
|
||||
def get_current_max_span(self):
|
||||
return self.decoder.get_current_max_span()
|
||||
|
||||
def get_current_avg_span(self):
|
||||
return self.decoder.get_current_avg_span()
|
||||
|
||||
|
||||
class AdaptiveSpanDecoder(FairseqIncrementalDecoder):
|
||||
def __init__(self, cfg, task):
|
||||
|
||||
super().__init__(task.target_dictionary)
|
||||
|
||||
self.config = cfg
|
||||
config = AdaptiveSpanSmallConfig(
|
||||
vocab_size=len(task.target_dictionary),
|
||||
d_model=cfg.d_model,
|
||||
n_head=cfg.n_head,
|
||||
d_inner=cfg.d_inner,
|
||||
n_layer=cfg.n_layer,
|
||||
attn_span=cfg.attn_span,
|
||||
dropout=cfg.dropout,
|
||||
emb_dropout=cfg.emb_dropout,
|
||||
adapt_span_ramp=cfg.adapt_span_ramp,
|
||||
adapt_span_init=cfg.adapt_span_init,
|
||||
aux_loss_scaler=cfg.aux_loss_scaler,
|
||||
adapt_span_layer=cfg.adapt_span_layer,
|
||||
)
|
||||
logger.info(config)
|
||||
self.model = AdaptiveSpanTransformerModel(**config.__dict__)
|
||||
|
||||
self._mems = None
|
||||
|
||||
def forward(
|
||||
self,
|
||||
src_tokens,
|
||||
incremental_state: Optional[Dict[str, List[torch.Tensor]]] = None,
|
||||
encoder_out=None,
|
||||
):
|
||||
bsz = src_tokens.size(0)
|
||||
if incremental_state is not None: # used during inference
|
||||
mems = self.get_incremental_state("mems")
|
||||
src_tokens = src_tokens[:, -1:] # only keep the most recent token
|
||||
else:
|
||||
mems = self._mems
|
||||
|
||||
if mems is None:
|
||||
# first time init
|
||||
mems = self.init_hid_cache(bsz)
|
||||
output = self.model(x=src_tokens, h_cache=mems,)
|
||||
if incremental_state is not None:
|
||||
self.set_incremental_state(incremental_state, "mems", output[1])
|
||||
else:
|
||||
self._mems = output[1]
|
||||
return (output[0],)
|
||||
|
||||
def max_positions(self):
|
||||
return self.config.attn_span
|
||||
|
||||
def init_hid_cache(self, batch_sz):
|
||||
hid = []
|
||||
for layer in self.model.layers:
|
||||
param = next(self.model.parameters())
|
||||
h = torch.zeros(
|
||||
batch_sz,
|
||||
layer.get_cache_size(),
|
||||
self.config.d_model,
|
||||
dtype=param.dtype,
|
||||
device=param.device,
|
||||
)
|
||||
hid.append(h)
|
||||
return hid
|
||||
|
||||
def get_aux_loss(self):
|
||||
return self.model.get_aux_loss()
|
||||
|
||||
def get_current_max_span(self):
|
||||
return self.model.get_current_max_span()
|
||||
|
||||
def get_current_avg_span(self):
|
||||
return self.model.get_current_avg_span()
|
||||
|
||||
def reorder_incremental_state(
|
||||
self,
|
||||
incremental_state: Dict[str, Dict[str, Optional[torch.Tensor]]],
|
||||
new_order: torch.Tensor,
|
||||
):
|
||||
"""Reorder incremental state.
|
||||
|
||||
This will be called when the order of the input has changed from the
|
||||
previous time step. A typical use case is beam search, where the input
|
||||
order changes between time steps based on the selection of beams.
|
||||
"""
|
||||
raise NotImplementedError("This is required for generation/beam search")
|
||||
# mems = self.get_incremental_state(incremental_state, "mems")
|
||||
# if mems is not None:
|
||||
# new_mems = [mems_i.index_select(1, new_order) for mems_i in mems]
|
||||
# self.set_incremental_state(incremental_state, "mems", new_mems)
|
||||
@@ -0,0 +1,280 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import logging
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import torch
|
||||
from fairseq import distributed_utils as dist_utils, utils
|
||||
from fairseq.data import (
|
||||
Dictionary,
|
||||
TokenBlockDataset,
|
||||
data_utils,
|
||||
iterators,
|
||||
)
|
||||
from fairseq.dataclass import FairseqDataclass
|
||||
from fairseq.tasks import FairseqTask, register_task
|
||||
from omegaconf import II
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class TruncatedBPTTLMConfig(FairseqDataclass):
|
||||
data: str = field(default="???", metadata={"help": "path to data directory"})
|
||||
tokens_per_sample: int = field(
|
||||
default=1024,
|
||||
metadata={"help": "max number of tokens per sequence"},
|
||||
)
|
||||
batch_size: int = II("dataset.batch_size")
|
||||
# Some models use *max_target_positions* to know how many positional
|
||||
# embeddings to learn. We use II(...) to make it default to
|
||||
# *tokens_per_sample*, but in principle there could be more positional
|
||||
# embeddings than tokens in a single batch. This may also be irrelevant for
|
||||
# custom model implementations.
|
||||
max_target_positions: int = II("task.tokens_per_sample")
|
||||
# these will be populated automatically if not provided
|
||||
data_parallel_rank: Optional[int] = None
|
||||
data_parallel_size: Optional[int] = None
|
||||
|
||||
|
||||
@register_task("truncated_bptt_lm", dataclass=TruncatedBPTTLMConfig)
|
||||
class TruncatedBPTTLMTask(FairseqTask):
|
||||
def __init__(self, cfg: TruncatedBPTTLMConfig):
|
||||
super().__init__(cfg)
|
||||
|
||||
if cfg.data_parallel_rank is None or cfg.data_parallel_size is None:
|
||||
if torch.distributed.is_initialized():
|
||||
cfg.data_parallel_rank = dist_utils.get_data_parallel_rank()
|
||||
cfg.data_parallel_size = dist_utils.get_data_parallel_world_size()
|
||||
else:
|
||||
cfg.data_parallel_rank = 0
|
||||
cfg.data_parallel_size = 1
|
||||
|
||||
# load the dictionary
|
||||
paths = utils.split_paths(cfg.data)
|
||||
assert len(paths) > 0
|
||||
self.dictionary = Dictionary.load(os.path.join(paths[0], "dict.txt"))
|
||||
logger.info("dictionary: {} types".format(len(self.dictionary)))
|
||||
|
||||
def load_dataset(self, split, epoch=1, combine=False, **kwargs):
|
||||
"""Load a given dataset split (e.g., train, valid, test)"""
|
||||
|
||||
# support sharded datasets
|
||||
paths = utils.split_paths(self.cfg.data)
|
||||
assert len(paths) > 0
|
||||
data_path = paths[(epoch - 1) % len(paths)]
|
||||
split_path = os.path.join(data_path, split)
|
||||
|
||||
# each element of *data* will be a tensorized line from the original
|
||||
# text dataset, similar to ``open(split_path).readlines()``
|
||||
data = data_utils.load_indexed_dataset(
|
||||
split_path, self.dictionary, combine=combine
|
||||
)
|
||||
if data is None:
|
||||
raise FileNotFoundError(
|
||||
"Dataset not found: {} ({})".format(split, split_path)
|
||||
)
|
||||
|
||||
# this is similar to ``data.view(-1).split(tokens_per_sample)``
|
||||
data = TokenBlockDataset(
|
||||
data,
|
||||
data.sizes,
|
||||
block_size=self.cfg.tokens_per_sample,
|
||||
pad=None, # unused
|
||||
eos=None, # unused
|
||||
break_mode="none",
|
||||
)
|
||||
|
||||
self.datasets[split] = TruncatedBPTTDataset(
|
||||
data=data,
|
||||
bsz_per_shard=self.cfg.batch_size,
|
||||
shard_id=self.cfg.data_parallel_rank,
|
||||
num_shards=self.cfg.data_parallel_size,
|
||||
)
|
||||
|
||||
def dataset(self, split):
|
||||
return self.datasets[split]
|
||||
|
||||
def get_batch_iterator(
|
||||
self, dataset, num_workers=0, epoch=1, data_buffer_size=0, **kwargs
|
||||
):
|
||||
return iterators.EpochBatchIterator(
|
||||
dataset=dataset,
|
||||
collate_fn=self._collate_fn,
|
||||
num_workers=num_workers,
|
||||
epoch=epoch,
|
||||
buffer_size=data_buffer_size,
|
||||
# we don't use the batching functionality from EpochBatchIterator;
|
||||
# instead every item in *dataset* is a whole batch
|
||||
batch_sampler=[[i] for i in range(len(dataset))],
|
||||
disable_shuffling=True,
|
||||
)
|
||||
|
||||
def _collate_fn(self, items: List[List[torch.Tensor]]):
|
||||
# we don't use fairseq's batching functionality, so we expect a single
|
||||
# Tensor of type List[torch.Tensor]
|
||||
assert len(items) == 1
|
||||
|
||||
# item will have shape B x T (the last batch may have length < T)
|
||||
id, item = items[0]
|
||||
item = data_utils.collate_tokens(item, pad_idx=self.source_dictionary.pad())
|
||||
B, T = item.size()
|
||||
|
||||
# shift item one position over and append a padding token for the target
|
||||
target = torch.nn.functional.pad(
|
||||
item[:, 1:], (0, 1, 0, 0), value=self.target_dictionary.pad()
|
||||
)
|
||||
|
||||
# fairseq expects batches to have the following structure
|
||||
return {
|
||||
"id": torch.tensor([id]*item.size(0)),
|
||||
"net_input": {
|
||||
"src_tokens": item,
|
||||
},
|
||||
"target": target,
|
||||
"nsentences": item.size(0),
|
||||
"ntokens": item.numel(),
|
||||
}
|
||||
|
||||
def build_dataset_for_inference(
|
||||
self, src_tokens: List[torch.Tensor], src_lengths: List[int], **kwargs
|
||||
) -> torch.utils.data.Dataset:
|
||||
eos = self.source_dictionary.eos()
|
||||
dataset = TokenBlockDataset(
|
||||
src_tokens,
|
||||
src_lengths,
|
||||
block_size=None, # ignored for "eos" break mode
|
||||
pad=self.source_dictionary.pad(),
|
||||
eos=eos,
|
||||
break_mode="eos",
|
||||
)
|
||||
|
||||
class Dataset(torch.utils.data.Dataset):
|
||||
def __getitem__(self, i):
|
||||
item = dataset[i]
|
||||
if item[-1] == eos:
|
||||
# remove eos to support generating with a prefix
|
||||
item = item[:-1]
|
||||
return (i, [item])
|
||||
|
||||
def __len__(self):
|
||||
return len(dataset)
|
||||
|
||||
return Dataset()
|
||||
|
||||
def inference_step(
|
||||
self, generator, models, sample, prefix_tokens=None, constraints=None
|
||||
):
|
||||
with torch.no_grad():
|
||||
if constraints is not None:
|
||||
raise NotImplementedError
|
||||
|
||||
# SequenceGenerator doesn't use *src_tokens* directly, we need to
|
||||
# pass the *prefix_tokens* argument instead.
|
||||
if prefix_tokens is None and sample["net_input"]["src_tokens"].nelement():
|
||||
prefix_tokens = sample["net_input"]["src_tokens"]
|
||||
|
||||
# begin generation with the end-of-sentence token
|
||||
bos_token = self.source_dictionary.eos()
|
||||
|
||||
return generator.generate(
|
||||
models, sample, prefix_tokens=prefix_tokens, bos_token=bos_token
|
||||
)
|
||||
|
||||
def eval_lm_dataloader(
|
||||
self,
|
||||
dataset,
|
||||
max_tokens: Optional[int] = 36000,
|
||||
batch_size: Optional[int] = None,
|
||||
max_positions: Optional[int] = None,
|
||||
num_shards: int = 1,
|
||||
shard_id: int = 0,
|
||||
num_workers: int = 1,
|
||||
data_buffer_size: int = 10,
|
||||
context_window: int = 0,
|
||||
):
|
||||
if context_window > 0:
|
||||
raise NotImplementedError(
|
||||
"Transformer-XL doesn't need --context-window, try "
|
||||
"--model-overrides '{\"mem_len\":42}' instead "
|
||||
)
|
||||
return self.get_batch_iterator(
|
||||
dataset=dataset,
|
||||
max_tokens=max_tokens,
|
||||
max_sentences=batch_size,
|
||||
max_positions=max_positions,
|
||||
ignore_invalid_inputs=True,
|
||||
num_shards=num_shards,
|
||||
shard_id=shard_id,
|
||||
num_workers=num_workers,
|
||||
data_buffer_size=data_buffer_size,
|
||||
).next_epoch_itr(shuffle=False)
|
||||
|
||||
@property
|
||||
def source_dictionary(self):
|
||||
return self.dictionary
|
||||
|
||||
@property
|
||||
def target_dictionary(self):
|
||||
return self.dictionary
|
||||
|
||||
|
||||
class TruncatedBPTTDataset(torch.utils.data.Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
data: List[torch.Tensor], # ordered list of items
|
||||
bsz_per_shard, # number of items processed per GPUs per forward
|
||||
shard_id, # current GPU ID
|
||||
num_shards, # number of GPUs
|
||||
):
|
||||
super().__init__()
|
||||
self.data = data
|
||||
|
||||
def batchify(data, bsz):
|
||||
# Work out how cleanly we can divide the dataset into bsz parts.
|
||||
nbatch = data.size(0) // bsz
|
||||
# Trim off any extra elements that wouldn't cleanly fit (remainders).
|
||||
data = data.narrow(0, 0, nbatch * bsz)
|
||||
# Evenly divide the data across the bsz batches.
|
||||
data = data.view(bsz, -1).contiguous()
|
||||
return data
|
||||
|
||||
# total number of sequences processed by all GPUs in each forward pass
|
||||
global_batch_size = bsz_per_shard * num_shards
|
||||
|
||||
"""
|
||||
With a 16 item dataset, bsz_per_shard=2 and num_shards=3,
|
||||
*indices* might look like:
|
||||
|
||||
indices = [[0, 1],
|
||||
[2, 3],
|
||||
[4, 5],
|
||||
[6, 7],
|
||||
[8, 9],
|
||||
[10, 11]]
|
||||
|
||||
The size of the TruncatedBPTTDataset instance will be 2,
|
||||
and shard 1 will see items:
|
||||
|
||||
[(0, [data[4], data[6]]),
|
||||
(1, [data[5], data[7]])]
|
||||
"""
|
||||
indices = batchify(torch.arange(len(data)), global_batch_size)
|
||||
assert indices.size(0) == global_batch_size
|
||||
|
||||
self.my_indices = indices[
|
||||
shard_id * bsz_per_shard : (shard_id + 1) * bsz_per_shard
|
||||
]
|
||||
assert self.my_indices.size(0) == bsz_per_shard
|
||||
|
||||
def __len__(self):
|
||||
return self.my_indices.size(1)
|
||||
|
||||
def __getitem__(self, i) -> Tuple[int, List[torch.Tensor]]:
|
||||
return (i, [self.data[idx] for idx in self.my_indices[:, i]])
|
||||
Reference in New Issue
Block a user