train.py 13.2 KB
Newer Older
Myle Ott's avatar
Myle Ott committed
1
#!/usr/bin/env python3 -u
Sergey Edunov's avatar
Sergey Edunov committed
2
3
4
5
6
7
8
# 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.

9
import collections
Myle Ott's avatar
Myle Ott committed
10
import itertools
11
12
13
import os
import math
import torch
Sergey Edunov's avatar
Sergey Edunov committed
14

Myle Ott's avatar
Myle Ott committed
15
from fairseq import data, distributed_utils, options, progress_bar, tasks, utils
16
17
18
from fairseq.fp16_trainer import FP16Trainer
from fairseq.trainer import Trainer
from fairseq.meters import AverageMeter, StopwatchMeter
Sergey Edunov's avatar
Sergey Edunov committed
19

Myle Ott's avatar
Myle Ott committed
20

Myle Ott's avatar
Myle Ott committed
21
def main(args):
22
23
    if args.max_tokens is None:
        args.max_tokens = 6000
24
25
26
27
28
29
30
    print(args)

    if not torch.cuda.is_available():
        raise NotImplementedError('Training on CPU is not supported')
    torch.cuda.set_device(args.device_id)
    torch.manual_seed(args.seed)

Myle Ott's avatar
Myle Ott committed
31
32
    # Setup task, e.g., translation, language modeling, etc.
    task = tasks.setup_task(args)
33

Myle Ott's avatar
Myle Ott committed
34
    # Load dataset splits
Alexei Baevski's avatar
Alexei Baevski committed
35
    load_dataset_splits(task, ['train', 'valid'])
36

Myle Ott's avatar
Myle Ott committed
37
38
39
    # Build model and criterion
    model = task.build_model(args)
    criterion = task.build_criterion(args)
40
    print('| model {}, criterion {}'.format(args.arch, criterion.__class__.__name__))
Myle Ott's avatar
Myle Ott committed
41
    print('| num. model params: {}'.format(sum(p.numel() for p in model.parameters())))
42
43
44

    # Build trainer
    if args.fp16:
45
46
47
        if torch.cuda.get_device_capability(0)[0] < 7:
            print('| WARNING: your device does NOT support faster training with --fp16,'
                  ' please switch to FP32 which is likely to be faster')
Myle Ott's avatar
Myle Ott committed
48
        trainer = FP16Trainer(args, task, model, criterion)
49
50
51
    else:
        if torch.cuda.get_device_capability(0)[0] >= 7:
            print('| NOTICE: your device may support faster training with --fp16')
Myle Ott's avatar
Myle Ott committed
52
        trainer = Trainer(args, task, model, criterion)
53
54
55
56
57
58
59
    print('| training on {} GPUs'.format(args.distributed_world_size))
    print('| max tokens per GPU = {} and max sentences per GPU = {}'.format(
        args.max_tokens,
        args.max_sentences,
    ))

    # Initialize dataloader
Myle Ott's avatar
Myle Ott committed
60
61
62
    max_positions = trainer.get_model().max_positions()
    epoch_itr = data.EpochBatchIterator(
        dataset=task.dataset(args.train_subset),
63
        max_tokens=args.max_tokens,
Myle Ott's avatar
Myle Ott committed
64
65
66
67
        max_sentences=args.max_sentences_valid,
        max_positions=max_positions,
        ignore_invalid_inputs=True,
        required_batch_size_multiple=8,
68
69
        seed=args.seed,
        num_shards=args.distributed_world_size,
Myle Ott's avatar
Myle Ott committed
70
        shard_id=args.distributed_rank,
71
72
73
    )

    # Load the latest checkpoint if one is available
Myle Ott's avatar
Myle Ott committed
74
    load_checkpoint(args, trainer, epoch_itr)
75
76

    # Send a dummy batch to warm the caching allocator
Myle Ott's avatar
Myle Ott committed
77
    dummy_batch = task.dataset('train').get_dummy_batch(args.max_tokens, max_positions)
78
79
80
81
82
83
84
85
    trainer.dummy_train_step(dummy_batch)

    # Train until the learning rate gets too small
    max_epoch = args.max_epoch or math.inf
    max_update = args.max_update or math.inf
    lr = trainer.get_lr()
    train_meter = StopwatchMeter()
    train_meter.start()
Myle Ott's avatar
Myle Ott committed
86
    valid_losses = [None]
87
    valid_subsets = args.valid_subset.split(',')
Myle Ott's avatar
Myle Ott committed
88
    while lr > args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates() < max_update:
89
        # train for one epoch
Myle Ott's avatar
Myle Ott committed
90
        train(args, trainer, task, epoch_itr)
91

Myle Ott's avatar
Myle Ott committed
92
93
        if epoch_itr.epoch % args.validate_interval == 0:
            valid_losses = validate(args, trainer, task, epoch_itr, valid_subsets)
94
95

        # only use first validation loss to update the learning rate
Myle Ott's avatar
Myle Ott committed
96
        lr = trainer.lr_step(epoch_itr.epoch, valid_losses[0])
97
98

        # save checkpoint
Myle Ott's avatar
Myle Ott committed
99
100
        if epoch_itr.epoch % args.save_interval == 0:
            save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
101
102
103
104
    train_meter.stop()
    print('| done training in {:.1f} seconds'.format(train_meter.sum))


Myle Ott's avatar
Myle Ott committed
105
def train(args, trainer, task, epoch_itr):
106
107
    """Train the model for one epoch."""

Myle Ott's avatar
Myle Ott committed
108
109
110
    # Initialize data iterator
    itr = epoch_itr.next_epoch_itr()
    progress = progress_bar.build_progress_bar(args, itr, epoch_itr.epoch, no_progress_bar='simple')
111
112

    # update parameters every N batches
Myle Ott's avatar
Myle Ott committed
113
114
    if epoch_itr.epoch <= len(args.update_freq):
        update_freq = args.update_freq[epoch_itr.epoch - 1]
115
116
117
118
    else:
        update_freq = args.update_freq[-1]

    extra_meters = collections.defaultdict(lambda: AverageMeter())
119
    first_valid = args.valid_subset.split(',')[0]
120
    max_update = args.max_update or math.inf
Myle Ott's avatar
Myle Ott committed
121
122
    num_batches = len(epoch_itr)
    for i, sample in enumerate(progress, start=epoch_itr.iterations_in_epoch):
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
        if i < num_batches - 1 and (i + 1) % update_freq > 0:
            # buffer updates according to --update-freq
            trainer.train_step(sample, update_params=False)
            continue
        else:
            log_output = trainer.train_step(sample, update_params=True)

        # log mid-epoch stats
        stats = get_training_stats(trainer)
        for k, v in log_output.items():
            if k in ['loss', 'nll_loss', 'sample_size']:
                continue  # these are already logged above
            if 'loss' in k:
                extra_meters[k].update(v, log_output['sample_size'])
            else:
                extra_meters[k].update(v)
            stats[k] = extra_meters[k].avg
        progress.log(stats)

        # ignore the first mini-batch in words-per-second calculation
        if i == 0:
            trainer.get_meter('wps').reset()

146
        num_updates = trainer.get_num_updates()
147
        if args.save_interval_updates > 0 and num_updates % args.save_interval_updates == 0 and num_updates > 0:
Myle Ott's avatar
Myle Ott committed
148
149
            valid_losses = validate(args, trainer, task, epoch_itr, [first_valid])
            save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
150
151

        if num_updates >= max_update:
152
153
154
155
156
157
158
159
            break

    # log end-of-epoch stats
    stats = get_training_stats(trainer)
    for k, meter in extra_meters.items():
        stats[k] = meter.avg
    progress.print(stats)

Myle Ott's avatar
Myle Ott committed
160
    # reset training meters
Sergey Edunov's avatar
Sergey Edunov committed
161
    for k in ['train_loss', 'train_nll_loss', 'wps', 'ups', 'wpb', 'bsz', 'clip', 'gnorm']:
Myle Ott's avatar
Myle Ott committed
162
163
164
165
        meter = trainer.get_meter(k)
        if meter is not None:
            meter.reset()

166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187

def get_training_stats(trainer):
    stats = collections.OrderedDict()
    stats['loss'] = '{:.3f}'.format(trainer.get_meter('train_loss').avg)
    if trainer.get_meter('train_nll_loss').count > 0:
        nll_loss = trainer.get_meter('train_nll_loss').avg
        stats['nll_loss'] = '{:.3f}'.format(nll_loss)
    else:
        nll_loss = trainer.get_meter('train_loss').avg
    stats['ppl'] = get_perplexity(nll_loss)
    stats['wps'] = round(trainer.get_meter('wps').avg)
    stats['ups'] = '{:.1f}'.format(trainer.get_meter('ups').avg)
    stats['wpb'] = round(trainer.get_meter('wpb').avg)
    stats['bsz'] = round(trainer.get_meter('bsz').avg)
    stats['num_updates'] = trainer.get_num_updates()
    stats['lr'] = trainer.get_lr()
    stats['gnorm'] = '{:.3f}'.format(trainer.get_meter('gnorm').avg)
    stats['clip'] = '{:.0%}'.format(trainer.get_meter('clip').avg)
    stats['oom'] = trainer.get_meter('oom').avg
    if trainer.get_meter('loss_scale') is not None:
        stats['loss_scale'] = '{:.3f}'.format(trainer.get_meter('loss_scale').avg)
    stats['wall'] = round(trainer.get_meter('wall').elapsed_time)
Myle Ott's avatar
Myle Ott committed
188
    stats['train_wall'] = round(trainer.get_meter('train_wall').sum)
189
190
191
    return stats


Myle Ott's avatar
Myle Ott committed
192
def validate(args, trainer, task, epoch_itr, subsets):
193
194
195
    """Evaluate the model on the validation set(s) and return the losses."""
    valid_losses = []
    for subset in subsets:
Myle Ott's avatar
Myle Ott committed
196
197
198
        # Initialize data iterator
        itr = data.EpochBatchIterator(
            dataset=task.dataset(subset),
199
200
            max_tokens=args.max_tokens,
            max_sentences=args.max_sentences_valid,
Myle Ott's avatar
Myle Ott committed
201
202
203
204
            max_positions=trainer.get_model().max_positions(),
            ignore_invalid_inputs=args.skip_invalid_size_inputs_valid_test,
            required_batch_size_multiple=8,
            seed=args.seed,
205
            num_shards=args.distributed_world_size,
Myle Ott's avatar
Myle Ott committed
206
207
            shard_id=args.distributed_rank,
        ).next_epoch_itr(shuffle=False)
208
        progress = progress_bar.build_progress_bar(
Myle Ott's avatar
Myle Ott committed
209
            args, itr, epoch_itr.epoch,
210
211
212
213
214
215
216
217
218
219
            prefix='valid on \'{}\' subset'.format(subset),
            no_progress_bar='simple'
        )

        # reset validation loss meters
        for k in ['valid_loss', 'valid_nll_loss']:
            meter = trainer.get_meter(k)
            if meter is not None:
                meter.reset()
        extra_meters = collections.defaultdict(lambda: AverageMeter())
Myle Ott's avatar
Myle Ott committed
220

221
222
223
224
225
226
227
        for sample in progress:
            log_output = trainer.valid_step(sample)

            for k, v in log_output.items():
                if k in ['loss', 'nll_loss', 'sample_size']:
                    continue
                extra_meters[k].update(v)
228

229
230
231
232
233
        # log validation stats
        stats = get_valid_stats(trainer)
        for k, meter in extra_meters.items():
            stats[k] = meter.avg
        progress.print(stats)
234

235
236
        valid_losses.append(stats['valid_loss'])
    return valid_losses
237
238
239
240
241
242
243
244


def get_valid_stats(trainer):
    stats = collections.OrderedDict()
    stats['valid_loss'] = trainer.get_meter('valid_loss').avg
    if trainer.get_meter('valid_nll_loss').count > 0:
        nll_loss = trainer.get_meter('valid_nll_loss').avg
        stats['valid_nll_loss'] = nll_loss
245
    else:
246
247
        nll_loss = trainer.get_meter('valid_loss').avg
    stats['valid_ppl'] = get_perplexity(nll_loss)
Myle Ott's avatar
Nits  
Myle Ott committed
248
249
250
    stats['num_updates'] = trainer.get_num_updates()
    if hasattr(save_checkpoint, 'best'):
        stats['best'] = min(save_checkpoint.best, stats['valid_loss'])
251
252
253
254
255
256
257
258
259
260
    return stats


def get_perplexity(loss):
    try:
        return '{:.2f}'.format(math.pow(2, loss))
    except OverflowError:
        return float('inf')


Myle Ott's avatar
Myle Ott committed
261
262
def save_checkpoint(args, trainer, epoch_itr, val_loss):
    if args.no_save or not distributed_utils.is_master(args):
263
        return
Myle Ott's avatar
Myle Ott committed
264
265
    epoch = epoch_itr.epoch
    end_of_epoch = epoch_itr.end_of_epoch()
266
267
268
269
    updates = trainer.get_num_updates()

    checkpoint_conds = collections.OrderedDict()
    checkpoint_conds['checkpoint{}.pt'.format(epoch)] = (
Alexei Baevski's avatar
Alexei Baevski committed
270
271
            end_of_epoch and not args.no_epoch_checkpoints and
            epoch % args.save_interval == 0
272
273
    )
    checkpoint_conds['checkpoint_{}_{}.pt'.format(epoch, updates)] = (
Alexei Baevski's avatar
Alexei Baevski committed
274
275
            not end_of_epoch and args.save_interval_updates > 0 and
            updates % args.save_interval_updates == 0
276
277
    )
    checkpoint_conds['checkpoint_best.pt'] = (
Alexei Baevski's avatar
Alexei Baevski committed
278
279
            val_loss is not None and
            (not hasattr(save_checkpoint, 'best') or val_loss < save_checkpoint.best)
280
281
282
    )
    checkpoint_conds['checkpoint_last.pt'] = True  # keep this last so that it's a symlink

Myle Ott's avatar
Myle Ott committed
283
284
285
    prev_best = getattr(save_checkpoint, 'best', val_loss)
    if val_loss is not None:
        save_checkpoint.best = min(val_loss, prev_best)
286
    extra_state = {
Myle Ott's avatar
Myle Ott committed
287
288
        'best': save_checkpoint.best,
        'train_iterator': epoch_itr.state_dict(),
289
290
291
        'val_loss': val_loss,
    }

292
293
    checkpoints = [os.path.join(args.save_dir, fn) for fn, cond in checkpoint_conds.items() if cond]
    if len(checkpoints) > 0:
294
295
        for cp in checkpoints:
            trainer.save_checkpoint(cp, extra_state)
296
297
298
299

    if not end_of_epoch and args.keep_interval_updates > 0:
        # remove old checkpoints; checkpoints are sorted in descending order
        checkpoints = utils.checkpoint_paths(args.save_dir, pattern=r'checkpoint_\d+_(\d+)\.pt')
300
301
        for old_chk in checkpoints[args.keep_interval_updates:]:
            os.remove(old_chk)
302
303


Myle Ott's avatar
Myle Ott committed
304
305
def load_checkpoint(args, trainer, epoch_itr):
    """Load a checkpoint and replay dataloader to match."""
306
307
308
    os.makedirs(args.save_dir, exist_ok=True)
    checkpoint_path = os.path.join(args.save_dir, args.restore_file)
    if os.path.isfile(checkpoint_path):
309
310
        extra_state = trainer.load_checkpoint(checkpoint_path, args.reset_optimizer, args.reset_lr_scheduler,
                                              eval(args.optimizer_overrides))
311
        if extra_state is not None:
Myle Ott's avatar
Myle Ott committed
312
313
314
315
316
            # replay train iterator to match checkpoint
            epoch_itr.load_state_dict(extra_state['train_iterator'])

            print('| loaded checkpoint {} (epoch {} @ {} updates)'.format(
                checkpoint_path, epoch_itr.epoch, trainer.get_num_updates()))
alexeib's avatar
alexeib committed
317

Myle Ott's avatar
Myle Ott committed
318
319
            trainer.lr_step(epoch_itr.epoch)
            trainer.lr_step_update(trainer.get_num_updates())
320
321
322
            if 'best' in extra_state:
                save_checkpoint.best = extra_state['best']

323

Alexei Baevski's avatar
Alexei Baevski committed
324
def load_dataset_splits(task, splits):
Myle Ott's avatar
Myle Ott committed
325
    for split in splits:
Alexei Baevski's avatar
Alexei Baevski committed
326
327
328
329
330
331
332
333
334
335
336
        if split == 'train':
            task.load_dataset(split, combine=True)
        else:
            for k in itertools.count():
                split_k = split + (str(k) if k > 0 else '')
                try:
                    task.load_dataset(split_k, combine=False)
                except FileNotFoundError as e:
                    if k > 0:
                        break
                    raise e
Sergey Edunov's avatar
Sergey Edunov committed
337

Myle Ott's avatar
Myle Ott committed
338

Sergey Edunov's avatar
Sergey Edunov committed
339
if __name__ == '__main__':
Myle Ott's avatar
Myle Ott committed
340
341
    parser = options.get_training_parser()
    args = options.parse_args_and_arch(parser)
342
343
344

    if args.distributed_port > 0 or args.distributed_init_method is not None:
        from distributed_train import main as distributed_main
345

346
347
348
        distributed_main(args)
    elif args.distributed_world_size > 1:
        from multiprocessing_train import main as multiprocessing_main
349

350
351
352
        multiprocessing_main(args)
    else:
        main(args)