train.py 14.9 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
# 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.
Myle Ott's avatar
Myle Ott committed
8
9
10
"""
Train a new model on one or across multiple GPUs.
"""
Sergey Edunov's avatar
Sergey Edunov committed
11

12
import collections
Myle Ott's avatar
Myle Ott committed
13
import itertools
14
15
import os
import math
16
17
import random

18
import torch
Sergey Edunov's avatar
Sergey Edunov committed
19

20
from fairseq import distributed_utils, options, progress_bar, tasks, utils
21
from fairseq.data import iterators
22
23
from fairseq.trainer import Trainer
from fairseq.meters import AverageMeter, StopwatchMeter
24
from fairseq.utils import import_user_module
Sergey Edunov's avatar
Sergey Edunov committed
25

Myle Ott's avatar
Myle Ott committed
26

Myle Ott's avatar
Myle Ott committed
27
def main(args):
28
29
    import_user_module(args)

30
31
    if args.max_tokens is None:
        args.max_tokens = 6000
32
33
    print(args)

Myle Ott's avatar
Myle Ott committed
34
35
    if torch.cuda.is_available() and not args.cpu:
        torch.cuda.set_device(args.device_id)
36
37
    torch.manual_seed(args.seed)

Myle Ott's avatar
Myle Ott committed
38
39
    # Setup task, e.g., translation, language modeling, etc.
    task = tasks.setup_task(args)
40

Myle Ott's avatar
Myle Ott committed
41
    # Load dataset splits
Alexei Baevski's avatar
Alexei Baevski committed
42
    load_dataset_splits(task, ['train', 'valid'])
43

Myle Ott's avatar
Myle Ott committed
44
45
46
    # Build model and criterion
    model = task.build_model(args)
    criterion = task.build_criterion(args)
47
    print('| model {}, criterion {}'.format(args.arch, criterion.__class__.__name__))
Myle Ott's avatar
Myle Ott committed
48
    print('| num. model params: {}'.format(sum(p.numel() for p in model.parameters())))
49

50
51
52
53
54
55
56
57
    # 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('train').get_dummy_batch(args.max_tokens, max_positions)
58
    oom_batch = task.dataset('train').get_dummy_batch(1, max_positions)
59

60
    # Build trainer
61
    trainer = Trainer(args, task, model, criterion, dummy_batch, oom_batch)
62
63
64
65
66
67
68
    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
69
    epoch_itr = task.get_batch_iterator(
Myle Ott's avatar
Myle Ott committed
70
        dataset=task.dataset(args.train_subset),
71
        max_tokens=args.max_tokens,
72
        max_sentences=args.max_sentences,
Myle Ott's avatar
Myle Ott committed
73
74
75
        max_positions=max_positions,
        ignore_invalid_inputs=True,
        required_batch_size_multiple=8,
76
77
        seed=args.seed,
        num_shards=args.distributed_world_size,
Myle Ott's avatar
Myle Ott committed
78
        shard_id=args.distributed_rank,
Myle Ott's avatar
Myle Ott committed
79
        num_workers=args.num_workers,
80
81
82
    )

    # Load the latest checkpoint if one is available
83
    if not load_checkpoint(args, trainer, epoch_itr):
84
        trainer.dummy_train_step([dummy_batch])
85
86
87
88
89
90
91

    # 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
92
    valid_losses = [None]
93
    valid_subsets = args.valid_subset.split(',')
Myle Ott's avatar
Myle Ott committed
94
    while lr > args.min_lr and epoch_itr.epoch < max_epoch and trainer.get_num_updates() < max_update:
95
        # train for one epoch
Myle Ott's avatar
Myle Ott committed
96
        train(args, trainer, task, epoch_itr)
97

Myle Ott's avatar
Myle Ott committed
98
99
        if epoch_itr.epoch % args.validate_interval == 0:
            valid_losses = validate(args, trainer, task, epoch_itr, valid_subsets)
100
101

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

        # save checkpoint
Myle Ott's avatar
Myle Ott committed
105
106
        if epoch_itr.epoch % args.save_interval == 0:
            save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
107
108
109
110
    train_meter.stop()
    print('| done training in {:.1f} seconds'.format(train_meter.sum))


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

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

120
    # Initialize data iterator
Myle Ott's avatar
Myle Ott committed
121
    itr = epoch_itr.next_epoch_itr(fix_batches_to_gpus=args.fix_batches_to_gpus)
122
123
124
125
126
    itr = iterators.GroupedIterator(itr, update_freq)
    progress = progress_bar.build_progress_bar(
        args, itr, epoch_itr.epoch, no_progress_bar='simple',
    )

127
    extra_meters = collections.defaultdict(lambda: AverageMeter())
128
    first_valid = args.valid_subset.split(',')[0]
129
    max_update = args.max_update or math.inf
130
131
132
    for i, samples in enumerate(progress, start=epoch_itr.iterations_in_epoch):
        log_output = trainer.train_step(samples)
        if log_output is None:
133
134
135
136
137
            continue

        # log mid-epoch stats
        stats = get_training_stats(trainer)
        for k, v in log_output.items():
138
            if k in ['loss', 'nll_loss', 'ntokens', 'nsentences', 'sample_size']:
139
140
141
142
143
144
145
146
147
148
149
150
                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()

151
        num_updates = trainer.get_num_updates()
152
        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
153
154
            valid_losses = validate(args, trainer, task, epoch_itr, [first_valid])
            save_checkpoint(args, trainer, epoch_itr, valid_losses[0])
155
156

        if num_updates >= max_update:
157
158
159
160
161
162
163
164
            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
165
    # reset training meters
166
167
168
    for k in [
        'train_loss', 'train_nll_loss', 'wps', 'ups', 'wpb', 'bsz', 'gnorm', 'clip',
    ]:
Myle Ott's avatar
Myle Ott committed
169
170
171
172
        meter = trainer.get_meter(k)
        if meter is not None:
            meter.reset()

173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194

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
195
    stats['train_wall'] = round(trainer.get_meter('train_wall').sum)
196
197
198
    return stats


Myle Ott's avatar
Myle Ott committed
199
def validate(args, trainer, task, epoch_itr, subsets):
200
201
202
    """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
203
        # Initialize data iterator
204
        itr = task.get_batch_iterator(
Myle Ott's avatar
Myle Ott committed
205
            dataset=task.dataset(subset),
206
207
            max_tokens=args.max_tokens,
            max_sentences=args.max_sentences_valid,
208
209
210
211
            max_positions=utils.resolve_max_positions(
                task.max_positions(),
                trainer.get_model().max_positions(),
            ),
Myle Ott's avatar
Myle Ott committed
212
213
214
            ignore_invalid_inputs=args.skip_invalid_size_inputs_valid_test,
            required_batch_size_multiple=8,
            seed=args.seed,
215
            num_shards=args.distributed_world_size,
Myle Ott's avatar
Myle Ott committed
216
            shard_id=args.distributed_rank,
Myle Ott's avatar
Myle Ott committed
217
            num_workers=args.num_workers,
Myle Ott's avatar
Myle Ott committed
218
        ).next_epoch_itr(shuffle=False)
219
        progress = progress_bar.build_progress_bar(
Myle Ott's avatar
Myle Ott committed
220
            args, itr, epoch_itr.epoch,
221
222
223
224
225
226
227
228
229
230
            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
231

232
233
234
235
        for sample in progress:
            log_output = trainer.valid_step(sample)

            for k, v in log_output.items():
236
                if k in ['loss', 'nll_loss', 'ntokens', 'nsentences', 'sample_size']:
237
238
                    continue
                extra_meters[k].update(v)
239

240
241
242
243
244
        # log validation stats
        stats = get_valid_stats(trainer)
        for k, meter in extra_meters.items():
            stats[k] = meter.avg
        progress.print(stats)
245

246
247
        valid_losses.append(stats['valid_loss'])
    return valid_losses
248
249
250
251
252
253
254
255


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
256
    else:
257
258
        nll_loss = trainer.get_meter('valid_loss').avg
    stats['valid_ppl'] = get_perplexity(nll_loss)
Myle Ott's avatar
Nits  
Myle Ott committed
259
260
261
    stats['num_updates'] = trainer.get_num_updates()
    if hasattr(save_checkpoint, 'best'):
        stats['best'] = min(save_checkpoint.best, stats['valid_loss'])
262
263
264
265
266
267
268
269
270
271
    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
272
273
def save_checkpoint(args, trainer, epoch_itr, val_loss):
    if args.no_save or not distributed_utils.is_master(args):
274
        return
Myle Ott's avatar
Myle Ott committed
275
276
    epoch = epoch_itr.epoch
    end_of_epoch = epoch_itr.end_of_epoch()
277
278
279
280
    updates = trainer.get_num_updates()

    checkpoint_conds = collections.OrderedDict()
    checkpoint_conds['checkpoint{}.pt'.format(epoch)] = (
Alexei Baevski's avatar
Alexei Baevski committed
281
282
            end_of_epoch and not args.no_epoch_checkpoints and
            epoch % args.save_interval == 0
283
284
    )
    checkpoint_conds['checkpoint_{}_{}.pt'.format(epoch, updates)] = (
Alexei Baevski's avatar
Alexei Baevski committed
285
286
            not end_of_epoch and args.save_interval_updates > 0 and
            updates % args.save_interval_updates == 0
287
288
    )
    checkpoint_conds['checkpoint_best.pt'] = (
Alexei Baevski's avatar
Alexei Baevski committed
289
290
            val_loss is not None and
            (not hasattr(save_checkpoint, 'best') or val_loss < save_checkpoint.best)
291
292
293
    )
    checkpoint_conds['checkpoint_last.pt'] = True  # keep this last so that it's a symlink

Myle Ott's avatar
Myle Ott committed
294
295
296
    prev_best = getattr(save_checkpoint, 'best', val_loss)
    if val_loss is not None:
        save_checkpoint.best = min(val_loss, prev_best)
297
    extra_state = {
Myle Ott's avatar
Myle Ott committed
298
        'train_iterator': epoch_itr.state_dict(),
299
300
        'val_loss': val_loss,
    }
Naman Goyal's avatar
Naman Goyal committed
301
302
    if hasattr(save_checkpoint, 'best'):
        extra_state.update({'best': save_checkpoint.best})
303

304
305
    checkpoints = [os.path.join(args.save_dir, fn) for fn, cond in checkpoint_conds.items() if cond]
    if len(checkpoints) > 0:
306
307
        for cp in checkpoints:
            trainer.save_checkpoint(cp, extra_state)
308
309
310
311

    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')
312
        for old_chk in checkpoints[args.keep_interval_updates:]:
Myle Ott's avatar
Myle Ott committed
313
314
315
316
317
318
319
320
321
            if os.path.lexists(old_chk):
                os.remove(old_chk)

    if args.keep_last_epochs > 0:
        # remove old epoch checkpoints; checkpoints are sorted in descending order
        checkpoints = utils.checkpoint_paths(args.save_dir, pattern=r'checkpoint\d+\.pt')
        for old_chk in checkpoints[args.keep_last_epochs:]:
            if os.path.lexists(old_chk):
                os.remove(old_chk)
322
323


Myle Ott's avatar
Myle Ott committed
324
325
def load_checkpoint(args, trainer, epoch_itr):
    """Load a checkpoint and replay dataloader to match."""
326
327
328
    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):
329
330
        extra_state = trainer.load_checkpoint(checkpoint_path, args.reset_optimizer, args.reset_lr_scheduler,
                                              eval(args.optimizer_overrides))
331
        if extra_state is not None:
Myle Ott's avatar
Myle Ott committed
332
333
334
335
336
            # 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
337

Myle Ott's avatar
Myle Ott committed
338
339
            trainer.lr_step(epoch_itr.epoch)
            trainer.lr_step_update(trainer.get_num_updates())
340
341
            if 'best' in extra_state:
                save_checkpoint.best = extra_state['best']
342
343
        return True
    return False
344

345

Alexei Baevski's avatar
Alexei Baevski committed
346
def load_dataset_splits(task, splits):
Myle Ott's avatar
Myle Ott committed
347
    for split in splits:
Alexei Baevski's avatar
Alexei Baevski committed
348
349
350
351
352
353
354
355
356
357
358
        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
359

Myle Ott's avatar
Myle Ott committed
360

Myle Ott's avatar
Myle Ott committed
361
362
363
364
365
366
367
368
369
370
def distributed_main(i, args):
    import socket
    args.device_id = i
    if args.distributed_rank is None:  # torch.multiprocessing.spawn
        args.distributed_rank = i
    args.distributed_rank = distributed_utils.distributed_init(args)
    print('| initialized host {} as rank {}'.format(socket.gethostname(), args.distributed_rank))
    main(args)


Sergey Edunov's avatar
Sergey Edunov committed
371
if __name__ == '__main__':
Myle Ott's avatar
Myle Ott committed
372
373
    parser = options.get_training_parser()
    args = options.parse_args_and_arch(parser)
374

Myle Ott's avatar
Myle Ott committed
375
376
    if args.distributed_init_method is None:
        distributed_utils.infer_init_method(args)
377

Myle Ott's avatar
Myle Ott committed
378
379
380
    if args.distributed_init_method is not None:
        # distributed training
        distributed_main(args.device_id, args)
381
    elif args.distributed_world_size > 1:
Myle Ott's avatar
Myle Ott committed
382
        # fallback for single node with multiple GPUs
383
384
        port = random.randint(10000, 20000)
        args.distributed_init_method = 'tcp://localhost:{port}'.format(port=port)
Myle Ott's avatar
Myle Ott committed
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
        args.distributed_rank = None  # set based on device id
        print(
            '''| NOTE: you may get better performance with:

            python -m torch.distributed.launch --nproc_per_node {ngpu} train.py {no_c10d}(...)
            '''.format(
                ngpu=args.distributed_world_size,
                no_c10d=(
                    '--ddp-backend=no_c10d ' if max(args.update_freq) > 1 and args.ddp_backend != 'no_c10d'
                    else ''
                ),
            )
        )
        torch.multiprocessing.spawn(
            fn=distributed_main,
            args=(args, ),
            nprocs=args.distributed_world_size,
        )
403
    else:
Myle Ott's avatar
Myle Ott committed
404
        # single GPU training
405
        main(args)