fairseq_task.py 7.23 KB
Newer Older
Myle Ott's avatar
Myle Ott committed
1
2
3
4
5
6
7
# Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the LICENSE file in
# the root directory of this source tree. An additional grant of patent rights
# can be found in the PATENTS file in the same directory.

8
from fairseq.data import data_utils, FairseqDataset, iterators
Peng-Jen Chen's avatar
Peng-Jen Chen committed
9
import torch
10

Myle Ott's avatar
Myle Ott committed
11
12
13

class FairseqTask(object):
    """
14
15
    Tasks store dictionaries and provide helpers for loading/iterating over
    Datasets, initializing the Model/Criterion and calculating the loss.
Myle Ott's avatar
Myle Ott committed
16
17
18
19
20
21
22
23
24
25
26
27
28
    """

    @staticmethod
    def add_args(parser):
        """Add task-specific arguments to the parser."""
        pass

    def __init__(self, args):
        self.args = args
        self.datasets = {}

    @classmethod
    def setup_task(cls, args, **kwargs):
Myle Ott's avatar
Myle Ott committed
29
30
31
32
33
        """Setup the task (e.g., load dictionaries).

        Args:
            args (argparse.Namespace): parsed command-line arguments
        """
34
        return cls(args)
Myle Ott's avatar
Myle Ott committed
35

Peng-Jen Chen's avatar
Peng-Jen Chen committed
36
    def load_dataset(self, split, combine=False, **kwargs):
Myle Ott's avatar
Myle Ott committed
37
38
39
40
41
        """Load a given dataset split.

        Args:
            split (str): name of the split (e.g., train, valid, test)
        """
Myle Ott's avatar
Myle Ott committed
42
43
44
        raise NotImplementedError

    def dataset(self, split):
Myle Ott's avatar
Myle Ott committed
45
46
47
48
49
50
51
52
53
        """
        Return a loaded dataset split.

        Args:
            split (str): name of the split (e.g., train, valid, test)

        Returns:
            a :class:`~fairseq.data.FairseqDataset` corresponding to *split*
        """
Myle Ott's avatar
Myle Ott committed
54
        from fairseq.data import FairseqDataset
Myle Ott's avatar
Myle Ott committed
55
56
57
58
59
60
        if split not in self.datasets:
            raise KeyError('Dataset not loaded: ' + split)
        if not isinstance(self.datasets[split], FairseqDataset):
            raise TypeError('Datasets are expected to be of type FairseqDataset')
        return self.datasets[split]

61
62
63
64
65
66
    def get_batch_iterator(
        self, dataset, max_tokens=None, max_sentences=None, max_positions=None,
        ignore_invalid_inputs=False, required_batch_size_multiple=1,
        seed=1, num_shards=1, shard_id=0,
    ):
        """
Myle Ott's avatar
Myle Ott committed
67
        Get an iterator that yields batches of data from the given dataset.
68
69

        Args:
Myle Ott's avatar
Myle Ott committed
70
            dataset (~fairseq.data.FairseqDataset): dataset to batch
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
            max_tokens (int, optional): max number of tokens in each batch.
                Default: ``None``
            max_sentences (int, optional): max number of sentences in each
                batch. Default: ``None``
            max_positions (optional): max sentence length supported by the
                model. Default: ``None``
            ignore_invalid_inputs (bool, optional): don't raise Exception for
                sentences that are too long. Default: ``False``
            required_batch_size_multiple (int, optional): require batch size to
                be a multiple of N. Default: ``1``
            seed (int, optional): seed for random number generator for
                reproducibility. Default: ``1``
            num_shards (int, optional): shard the data iterator into N
                shards. Default: ``1``
            shard_id (int, optional): which shard of the data iterator to
                return. Default: ``0``

        Returns:
Myle Ott's avatar
Myle Ott committed
89
90
            ~fairseq.iterators.EpochBatchIterator: a batched iterator over the
                given dataset split
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
        """
        assert isinstance(dataset, FairseqDataset)

        # get indices ordered by example size
        with data_utils.numpy_seed(seed):
            indices = dataset.ordered_indices()

        # filter examples that are too large
        indices = data_utils.filter_by_size(
            indices, dataset.size, max_positions, raise_exception=(not ignore_invalid_inputs),
        )

        # create mini-batches with given size constraints
        batch_sampler = data_utils.batch_by_size(
            indices, dataset.num_tokens, max_tokens=max_tokens, max_sentences=max_sentences,
            required_batch_size_multiple=required_batch_size_multiple,
        )

        # return a reusable, sharded iterator
110
        return iterators.EpochBatchIterator(
111
            dataset=dataset,
112
            collate_fn=dataset.collater,
113
114
115
116
117
118
            batch_sampler=batch_sampler,
            seed=seed,
            num_shards=num_shards,
            shard_id=shard_id,
        )

Myle Ott's avatar
Myle Ott committed
119
    def build_model(self, args):
Myle Ott's avatar
Myle Ott committed
120
121
122
123
124
125
126
127
128
129
        """
        Build the :class:`~fairseq.models.BaseFairseqModel` instance for this
        task.

        Args:
            args (argparse.Namespace): parsed command-line arguments

        Returns:
            a :class:`~fairseq.models.BaseFairseqModel` instance
        """
Myle Ott's avatar
Myle Ott committed
130
        from fairseq import models
Myle Ott's avatar
Myle Ott committed
131
132
133
        return models.build_model(args, self)

    def build_criterion(self, args):
Myle Ott's avatar
Myle Ott committed
134
135
136
137
138
139
140
141
142
143
        """
        Build the :class:`~fairseq.criterions.FairseqCriterion` instance for
        this task.

        Args:
            args (argparse.Namespace): parsed command-line arguments

        Returns:
            a :class:`~fairseq.criterions.FairseqCriterion` instance
        """
Myle Ott's avatar
Myle Ott committed
144
        from fairseq import criterions
Myle Ott's avatar
Myle Ott committed
145
146
        return criterions.build_criterion(args, self)

Peng-Jen Chen's avatar
Peng-Jen Chen committed
147
    def train_step(self, sample, model, criterion, optimizer, ignore_grad=False):
Myle Ott's avatar
Myle Ott committed
148
        """
Peng-Jen Chen's avatar
Peng-Jen Chen committed
149
150
        Do forward and backward, and return the loss as computed by *criterion*
        for the given *model* and *sample*.
Myle Ott's avatar
Myle Ott committed
151
152
153
154

        Args:
            sample (dict): the mini-batch. The format is defined by the
                :class:`~fairseq.data.FairseqDataset`.
Peng-Jen Chen's avatar
Peng-Jen Chen committed
155
156
157
158
159
160
161
162
163
164
165
            model (~fairseq.models.BaseFairseqModel): the model
            criterion (~fairseq.criterions.FairseqCriterion): the criterion
            optimizer (~fairseq.optim.FairseqOptimizer): the optimizer
            ignore_grad (bool): multiply loss by 0 if this is set to True

        Returns:
            tuple:
                - the loss
                - the sample size, which is used as the denominator for the
                  gradient
                - logging outputs to display while training
Myle Ott's avatar
Myle Ott committed
166
        """
Peng-Jen Chen's avatar
Peng-Jen Chen committed
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
        model.train()

        loss, sample_size, logging_output = criterion(model, sample)
        if ignore_grad:
            loss *= 0
        optimizer.backward(loss)
        return loss, sample_size, logging_output

    def valid_step(self, sample, model, criterion):
        model.eval()
        with torch.no_grad():
            loss, sample_size, logging_output = criterion(model, sample)
        return loss, sample_size, logging_output

    def init_logging_output(self, sample):
        return {
            'ntokens': sample['ntokens'] if sample is not None else 0,
            'nsentences': sample['target'].size(0) if sample is not None else 0,
        }

    def grad_denom(self, sample_sizes, criterion):
        return criterion.__class__.grad_denom(sample_sizes)

    def aggregate_logging_outputs(self, logging_outputs, criterion):
Myle Ott's avatar
Myle Ott committed
191
        return criterion._aggregate_logging_outputs(logging_outputs)
Myle Ott's avatar
Myle Ott committed
192

193
    def max_positions(self):
Myle Ott's avatar
Myle Ott committed
194
        """Return the max input length allowed by the task."""
195
196
        return None

Myle Ott's avatar
Myle Ott committed
197
198
    @property
    def source_dictionary(self):
Myle Ott's avatar
Myle Ott committed
199
200
        """Return the source :class:`~fairseq.data.Dictionary` (if applicable
        for this task)."""
Myle Ott's avatar
Myle Ott committed
201
202
203
204
        raise NotImplementedError

    @property
    def target_dictionary(self):
Myle Ott's avatar
Myle Ott committed
205
206
        """Return the target :class:`~fairseq.data.Dictionary` (if applicable
        for this task)."""
Myle Ott's avatar
Myle Ott committed
207
        raise NotImplementedError