mirror of
https://github.com/facebookresearch/fairseq.git
synced 2024-09-11 17:25:31 +03:00
Deprecate dummy_batch (#699)
Summary: It was tedious defining these, let's try just taking the first batch lazily instead. Pull Request resolved: https://github.com/pytorch/fairseq/pull/699 Differential Revision: D15188266 Pulled By: myleott fbshipit-source-id: a4c9f7ee3111278faaffa8a22ba91ed5f50e143d
This commit is contained in:
parent
7a5996fdc7
commit
fc1a19a38d
@ -10,7 +10,6 @@ import torch
|
|||||||
from fairseq import utils
|
from fairseq import utils
|
||||||
|
|
||||||
from . import FairseqDataset
|
from . import FairseqDataset
|
||||||
from .language_pair_dataset import collate as language_pair_collate, generate_dummy_batch
|
|
||||||
|
|
||||||
|
|
||||||
def backtranslate_samples(samples, collate_fn, generate_fn, cuda=True):
|
def backtranslate_samples(samples, collate_fn, generate_fn, cuda=True):
|
||||||
@ -141,19 +140,6 @@ class BacktranslationDataset(FairseqDataset):
|
|||||||
)
|
)
|
||||||
return self.output_collater(samples)
|
return self.output_collater(samples)
|
||||||
|
|
||||||
def get_dummy_batch(self, num_tokens, max_positions):
|
|
||||||
"""Just use the tgt dataset get_dummy_batch"""
|
|
||||||
def collate_fn(samples):
|
|
||||||
return language_pair_collate(
|
|
||||||
samples, pad_idx=self.src_dict.pad(), eos_idx=self.src_dict.eos(),
|
|
||||||
input_feeding=True,
|
|
||||||
)
|
|
||||||
dummy_batch = generate_dummy_batch(
|
|
||||||
num_tokens, collate_fn,
|
|
||||||
self.src_dict, tgt_dict=self.tgt_dict)
|
|
||||||
dummy_batch['is_dummy'] = True
|
|
||||||
return dummy_batch
|
|
||||||
|
|
||||||
def num_tokens(self, index):
|
def num_tokens(self, index):
|
||||||
"""Just use the tgt dataset num_tokens"""
|
"""Just use the tgt dataset num_tokens"""
|
||||||
return self.tgt_dataset.num_tokens(index)
|
return self.tgt_dataset.num_tokens(index)
|
||||||
|
@ -28,10 +28,6 @@ class FairseqDataset(torch.utils.data.Dataset):
|
|||||||
"""
|
"""
|
||||||
raise NotImplementedError
|
raise NotImplementedError
|
||||||
|
|
||||||
def get_dummy_batch(self, num_tokens, max_positions):
|
|
||||||
"""Return a dummy batch with a given number of tokens."""
|
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
def num_tokens(self, index):
|
def num_tokens(self, index):
|
||||||
"""Return the number of tokens in a sample. This value is used to
|
"""Return the number of tokens in a sample. This value is used to
|
||||||
enforce ``--max-tokens`` during batching."""
|
enforce ``--max-tokens`` during batching."""
|
||||||
|
@ -68,19 +68,6 @@ def collate(
|
|||||||
return batch
|
return batch
|
||||||
|
|
||||||
|
|
||||||
def generate_dummy_batch(num_tokens, collate_fn, src_dict, src_len=128, tgt_dict=None, tgt_len=128):
|
|
||||||
"""Return a dummy batch with a given number of tokens."""
|
|
||||||
bsz = num_tokens // max(src_len, tgt_len)
|
|
||||||
return collate_fn([
|
|
||||||
{
|
|
||||||
'id': i,
|
|
||||||
'source': src_dict.dummy_sentence(src_len),
|
|
||||||
'target': tgt_dict.dummy_sentence(tgt_len) if tgt_dict is not None else None,
|
|
||||||
}
|
|
||||||
for i in range(bsz)
|
|
||||||
])
|
|
||||||
|
|
||||||
|
|
||||||
class LanguagePairDataset(FairseqDataset):
|
class LanguagePairDataset(FairseqDataset):
|
||||||
"""
|
"""
|
||||||
A pair of torch.utils.data.Datasets.
|
A pair of torch.utils.data.Datasets.
|
||||||
@ -198,15 +185,6 @@ class LanguagePairDataset(FairseqDataset):
|
|||||||
input_feeding=self.input_feeding,
|
input_feeding=self.input_feeding,
|
||||||
)
|
)
|
||||||
|
|
||||||
def get_dummy_batch(self, num_tokens, max_positions, src_len=128, tgt_len=128):
|
|
||||||
"""Return a dummy batch with a given number of tokens."""
|
|
||||||
src_len, tgt_len = utils.resolve_max_positions(
|
|
||||||
(src_len, tgt_len),
|
|
||||||
max_positions,
|
|
||||||
(self.max_source_positions, self.max_target_positions),
|
|
||||||
)
|
|
||||||
return generate_dummy_batch(num_tokens, self.collater, self.src_dict, src_len, self.tgt_dict, tgt_len)
|
|
||||||
|
|
||||||
def num_tokens(self, index):
|
def num_tokens(self, index):
|
||||||
"""Return the number of tokens in a sample. This value is used to
|
"""Return the number of tokens in a sample. This value is used to
|
||||||
enforce ``--max-tokens`` during batching."""
|
enforce ``--max-tokens`` during batching."""
|
||||||
|
@ -62,9 +62,6 @@ class LMContextWindowDataset(FairseqDataset):
|
|||||||
|
|
||||||
return sample
|
return sample
|
||||||
|
|
||||||
def get_dummy_batch(self, *args, **kwargs):
|
|
||||||
return self.dataset.get_dummy_batch(*args, **kwargs)
|
|
||||||
|
|
||||||
def num_tokens(self, index):
|
def num_tokens(self, index):
|
||||||
return self.dataset.num_tokens(index)
|
return self.dataset.num_tokens(index)
|
||||||
|
|
||||||
|
@ -294,33 +294,6 @@ class MaskedLMDataset(FairseqDataset):
|
|||||||
"""
|
"""
|
||||||
return self._collate(samples, self.vocab.pad(), self.vocab.eos())
|
return self._collate(samples, self.vocab.pad(), self.vocab.eos())
|
||||||
|
|
||||||
def get_dummy_batch(
|
|
||||||
self,
|
|
||||||
num_tokens: int,
|
|
||||||
max_positions: Union[float, int],
|
|
||||||
tgt_len: int = 12
|
|
||||||
):
|
|
||||||
"""
|
|
||||||
Return a dummy batch with a given number of tokens.
|
|
||||||
"""
|
|
||||||
if isinstance(max_positions, float) or isinstance(max_positions, int):
|
|
||||||
tgt_len = min(tgt_len, max_positions)
|
|
||||||
source = self.vocab.dummy_sentence(tgt_len)
|
|
||||||
sentence_target = 0
|
|
||||||
bsz = num_tokens // tgt_len
|
|
||||||
|
|
||||||
return self.collater(
|
|
||||||
[
|
|
||||||
{
|
|
||||||
"id": i,
|
|
||||||
"block_one": source,
|
|
||||||
"block_two": source if self.has_pairs else None,
|
|
||||||
"sentence_target": sentence_target if self.has_pairs else None,
|
|
||||||
}
|
|
||||||
for i in range(bsz)
|
|
||||||
]
|
|
||||||
)
|
|
||||||
|
|
||||||
def num_tokens(
|
def num_tokens(
|
||||||
self,
|
self,
|
||||||
index: int
|
index: int
|
||||||
@ -358,4 +331,4 @@ class MaskedLMDataset(FairseqDataset):
|
|||||||
return getattr(self.dataset, "supports_prefetch", False)
|
return getattr(self.dataset, "supports_prefetch", False)
|
||||||
|
|
||||||
def prefetch(self, indices):
|
def prefetch(self, indices):
|
||||||
self.dataset.prefetch(indices)
|
self.dataset.prefetch(indices)
|
||||||
|
@ -174,21 +174,6 @@ class MonolingualDataset(FairseqDataset):
|
|||||||
"""
|
"""
|
||||||
return collate(samples, self.vocab.pad(), self.vocab.eos())
|
return collate(samples, self.vocab.pad(), self.vocab.eos())
|
||||||
|
|
||||||
def get_dummy_batch(self, num_tokens, max_positions, tgt_len=128):
|
|
||||||
"""Return a dummy batch with a given number of tokens."""
|
|
||||||
if isinstance(max_positions, float) or isinstance(max_positions, int):
|
|
||||||
tgt_len = min(tgt_len, max_positions)
|
|
||||||
bsz = max(num_tokens // tgt_len, 1)
|
|
||||||
target = self.vocab.dummy_sentence(tgt_len + 2)
|
|
||||||
source, past_target, future_target = target[1:-1], target[2:], target[:-2]
|
|
||||||
source, target = self._make_source_target(source, past_target, future_target)
|
|
||||||
source, target = self._maybe_add_bos(source, target)
|
|
||||||
|
|
||||||
return self.collater([
|
|
||||||
{'id': i, 'source': source, 'target': target}
|
|
||||||
for i in range(bsz)
|
|
||||||
])
|
|
||||||
|
|
||||||
def num_tokens(self, index):
|
def num_tokens(self, index):
|
||||||
"""Return the number of tokens in a sample. This value is used to
|
"""Return the number of tokens in a sample. This value is used to
|
||||||
enforce ``--max-tokens`` during batching."""
|
enforce ``--max-tokens`` during batching."""
|
||||||
|
@ -116,15 +116,6 @@ class MultiCorpusSampledDataset(FairseqDataset):
|
|||||||
selected_samples = [sample[selected_key] for sample in samples]
|
selected_samples = [sample[selected_key] for sample in samples]
|
||||||
return self.datasets[selected_key].collater(selected_samples)
|
return self.datasets[selected_key].collater(selected_samples)
|
||||||
|
|
||||||
def get_dummy_batch(self, num_tokens: int, max_positions: int):
|
|
||||||
"""
|
|
||||||
Return a dummy batch with a given number of tokens. Assumes that the
|
|
||||||
max_positions specified is the same for all underlying datasets.
|
|
||||||
"""
|
|
||||||
return self.datasets[self.default_key].get_dummy_batch(
|
|
||||||
num_tokens, max_positions
|
|
||||||
)
|
|
||||||
|
|
||||||
def num_tokens(self, index: int):
|
def num_tokens(self, index: int):
|
||||||
"""
|
"""
|
||||||
Return an example's length (number of tokens), used for batching. Here
|
Return an example's length (number of tokens), used for batching. Here
|
||||||
|
@ -72,17 +72,6 @@ class RoundRobinZipDatasets(FairseqDataset):
|
|||||||
# at evaluation time it's useful to pass-through batches from a single key
|
# at evaluation time it's useful to pass-through batches from a single key
|
||||||
return self.datasets[self.eval_key].collater(samples)
|
return self.datasets[self.eval_key].collater(samples)
|
||||||
|
|
||||||
def get_dummy_batch(self, max_tokens, max_positions):
|
|
||||||
if self.eval_key is None:
|
|
||||||
# TODO should max_tokens be used independently for each batch like this?
|
|
||||||
return OrderedDict([
|
|
||||||
(key, dataset.get_dummy_batch(max_tokens, max_positions[key]))
|
|
||||||
for key, dataset in self.datasets.items()
|
|
||||||
])
|
|
||||||
else:
|
|
||||||
# at evaluation time it's useful to return a single batch directly
|
|
||||||
return self.datasets[self.eval_key].get_dummy_batch(max_tokens, max_positions[self.eval_key])
|
|
||||||
|
|
||||||
def num_tokens(self, index):
|
def num_tokens(self, index):
|
||||||
"""Return an example's length (number of tokens), used for batching."""
|
"""Return an example's length (number of tokens), used for batching."""
|
||||||
# TODO make it configurable whether to use max() or sum() here
|
# TODO make it configurable whether to use max() or sum() here
|
||||||
|
@ -96,9 +96,6 @@ class TransformEosDataset(FairseqDataset):
|
|||||||
samples = list(map(transform, samples))
|
samples = list(map(transform, samples))
|
||||||
return self.dataset.collater(samples)
|
return self.dataset.collater(samples)
|
||||||
|
|
||||||
def get_dummy_batch(self, *args, **kwargs):
|
|
||||||
return self.dataset.get_dummy_batch(*args, **kwargs)
|
|
||||||
|
|
||||||
def num_tokens(self, index):
|
def num_tokens(self, index):
|
||||||
return self.dataset.num_tokens(index)
|
return self.dataset.num_tokens(index)
|
||||||
|
|
||||||
|
@ -60,9 +60,6 @@ class TransformEosLangPairDataset(FairseqDataset):
|
|||||||
|
|
||||||
return samples
|
return samples
|
||||||
|
|
||||||
def get_dummy_batch(self, *args, **kwargs):
|
|
||||||
return self.dataset.get_dummy_batch(*args, **kwargs)
|
|
||||||
|
|
||||||
def num_tokens(self, index):
|
def num_tokens(self, index):
|
||||||
return self.dataset.num_tokens(index)
|
return self.dataset.num_tokens(index)
|
||||||
|
|
||||||
|
@ -32,7 +32,7 @@ class Trainer(object):
|
|||||||
communication of the gradients across workers.
|
communication of the gradients across workers.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(self, args, task, model, criterion, dummy_batch, oom_batch=None):
|
def __init__(self, args, task, model, criterion, dummy_batch=None, oom_batch=None):
|
||||||
self.args = args
|
self.args = args
|
||||||
self.task = task
|
self.task = task
|
||||||
|
|
||||||
@ -47,7 +47,7 @@ class Trainer(object):
|
|||||||
self._model = self._model.cuda()
|
self._model = self._model.cuda()
|
||||||
|
|
||||||
self._dummy_batch = dummy_batch
|
self._dummy_batch = dummy_batch
|
||||||
self._oom_batch = oom_batch
|
self._oom_batch = oom_batch or dummy_batch
|
||||||
|
|
||||||
self._lr_scheduler = None
|
self._lr_scheduler = None
|
||||||
self._num_updates = 0
|
self._num_updates = 0
|
||||||
@ -177,6 +177,9 @@ class Trainer(object):
|
|||||||
|
|
||||||
def train_step(self, samples, dummy_batch=False, raise_oom=False):
|
def train_step(self, samples, dummy_batch=False, raise_oom=False):
|
||||||
"""Do forward, backward and parameter update."""
|
"""Do forward, backward and parameter update."""
|
||||||
|
if self._dummy_batch is None:
|
||||||
|
self._dummy_batch = samples[0]
|
||||||
|
|
||||||
self._set_seed()
|
self._set_seed()
|
||||||
self.model.train()
|
self.model.train()
|
||||||
self.criterion.train()
|
self.criterion.train()
|
||||||
|
20
train.py
20
train.py
@ -50,18 +50,8 @@ def main(args):
|
|||||||
sum(p.numel() for p in model.parameters() if p.requires_grad),
|
sum(p.numel() for p in model.parameters() if p.requires_grad),
|
||||||
))
|
))
|
||||||
|
|
||||||
# Make a dummy batch to (i) warm the caching allocator and (ii) as a
|
|
||||||
# placeholder DistributedDataParallel when there's an uneven number of
|
|
||||||
# batches per worker.
|
|
||||||
max_positions = utils.resolve_max_positions(
|
|
||||||
task.max_positions(),
|
|
||||||
model.max_positions(),
|
|
||||||
)
|
|
||||||
dummy_batch = task.dataset(args.train_subset).get_dummy_batch(args.max_tokens, max_positions)
|
|
||||||
oom_batch = task.dataset(args.train_subset).get_dummy_batch(1, max_positions)
|
|
||||||
|
|
||||||
# Build trainer
|
# Build trainer
|
||||||
trainer = Trainer(args, task, model, criterion, dummy_batch, oom_batch)
|
trainer = Trainer(args, task, model, criterion)
|
||||||
print('| training on {} GPUs'.format(args.distributed_world_size))
|
print('| training on {} GPUs'.format(args.distributed_world_size))
|
||||||
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
|
print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
|
||||||
args.max_tokens,
|
args.max_tokens,
|
||||||
@ -73,7 +63,10 @@ def main(args):
|
|||||||
dataset=task.dataset(args.train_subset),
|
dataset=task.dataset(args.train_subset),
|
||||||
max_tokens=args.max_tokens,
|
max_tokens=args.max_tokens,
|
||||||
max_sentences=args.max_sentences,
|
max_sentences=args.max_sentences,
|
||||||
max_positions=max_positions,
|
max_positions=utils.resolve_max_positions(
|
||||||
|
task.max_positions(),
|
||||||
|
model.max_positions(),
|
||||||
|
),
|
||||||
ignore_invalid_inputs=True,
|
ignore_invalid_inputs=True,
|
||||||
required_batch_size_multiple=args.required_batch_size_multiple,
|
required_batch_size_multiple=args.required_batch_size_multiple,
|
||||||
seed=args.seed,
|
seed=args.seed,
|
||||||
@ -83,8 +76,7 @@ def main(args):
|
|||||||
)
|
)
|
||||||
|
|
||||||
# Load the latest checkpoint if one is available
|
# Load the latest checkpoint if one is available
|
||||||
if not load_checkpoint(args, trainer, epoch_itr):
|
load_checkpoint(args, trainer, epoch_itr)
|
||||||
trainer.dummy_train_step([dummy_batch])
|
|
||||||
|
|
||||||
# Train until the learning rate gets too small
|
# Train until the learning rate gets too small
|
||||||
max_epoch = args.max_epoch or math.inf
|
max_epoch = args.max_epoch or math.inf
|
||||||
|
Loading…
Reference in New Issue
Block a user