train.py 13.5 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

15
from fairseq import 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
60
61
62
63
64
    max_positions = utils.resolve_max_positions(
        task.max_positions(),
        trainer.get_model().max_positions(),
    )
    epoch_itr = task.get_batch_iterator(
Myle Ott's avatar
Myle Ott committed
65
        dataset=task.dataset(args.train_subset),
66
        max_tokens=args.max_tokens,
67
        max_sentences=args.max_sentences,
Myle Ott's avatar
Myle Ott committed
68
69
70
        max_positions=max_positions,
        ignore_invalid_inputs=True,
        required_batch_size_multiple=8,
71
72
        seed=args.seed,
        num_shards=args.distributed_world_size,
Myle Ott's avatar
Myle Ott committed
73
        shard_id=args.distributed_rank,
74
75
76
    )

    # Load the latest checkpoint if one is available
77
78
79
80
    if not load_checkpoint(args, trainer, epoch_itr):
        # Send a dummy batch to warm the caching allocator
        dummy_batch = task.dataset('train').get_dummy_batch(args.max_tokens, max_positions)
        trainer.dummy_train_step(dummy_batch)
81
82
83
84
85
86
87

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

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

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

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


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

Myle Ott's avatar
Myle Ott committed
110
111
112
    # Initialize data iterator
    itr = epoch_itr.next_epoch_itr()
    progress = progress_bar.build_progress_bar(args, itr, epoch_itr.epoch, no_progress_bar='simple')
113
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
120
    else:
        update_freq = args.update_freq[-1]

    extra_meters = collections.defaultdict(lambda: AverageMeter())
121
    first_valid = args.valid_subset.split(',')[0]
122
    max_update = args.max_update or math.inf
Myle Ott's avatar
Myle Ott committed
123
124
    num_batches = len(epoch_itr)
    for i, sample in enumerate(progress, start=epoch_itr.iterations_in_epoch):
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
        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()

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

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

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

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
190
    stats['train_wall'] = round(trainer.get_meter('train_wall').sum)
191
192
193
    return stats


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

226
227
228
229
230
231
232
        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)
233

234
235
236
237
238
        # log validation stats
        stats = get_valid_stats(trainer)
        for k, meter in extra_meters.items():
            stats[k] = meter.avg
        progress.print(stats)
239

240
241
        valid_losses.append(stats['valid_loss'])
    return valid_losses
242
243
244
245
246
247
248
249


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

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

Myle Ott's avatar
Myle Ott committed
288
289
290
    prev_best = getattr(save_checkpoint, 'best', val_loss)
    if val_loss is not None:
        save_checkpoint.best = min(val_loss, prev_best)
291
    extra_state = {
Myle Ott's avatar
Myle Ott committed
292
293
        'best': save_checkpoint.best,
        'train_iterator': epoch_itr.state_dict(),
294
295
296
        'val_loss': val_loss,
    }

297
298
    checkpoints = [os.path.join(args.save_dir, fn) for fn, cond in checkpoint_conds.items() if cond]
    if len(checkpoints) > 0:
299
300
        for cp in checkpoints:
            trainer.save_checkpoint(cp, extra_state)
301
302
303
304

    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')
305
306
        for old_chk in checkpoints[args.keep_interval_updates:]:
            os.remove(old_chk)
307
308


Myle Ott's avatar
Myle Ott committed
309
310
def load_checkpoint(args, trainer, epoch_itr):
    """Load a checkpoint and replay dataloader to match."""
311
312
313
    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):
314
315
        extra_state = trainer.load_checkpoint(checkpoint_path, args.reset_optimizer, args.reset_lr_scheduler,
                                              eval(args.optimizer_overrides))
316
        if extra_state is not None:
Myle Ott's avatar
Myle Ott committed
317
318
319
320
321
            # 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
322

Myle Ott's avatar
Myle Ott committed
323
324
            trainer.lr_step(epoch_itr.epoch)
            trainer.lr_step_update(trainer.get_num_updates())
325
326
            if 'best' in extra_state:
                save_checkpoint.best = extra_state['best']
327
328
        return True
    return False
329

330

Alexei Baevski's avatar
Alexei Baevski committed
331
def load_dataset_splits(task, splits):
Myle Ott's avatar
Myle Ott committed
332
    for split in splits:
Alexei Baevski's avatar
Alexei Baevski committed
333
334
335
336
337
338
339
340
341
342
343
        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
344

Myle Ott's avatar
Myle Ott committed
345

Sergey Edunov's avatar
Sergey Edunov committed
346
if __name__ == '__main__':
Myle Ott's avatar
Myle Ott committed
347
348
    parser = options.get_training_parser()
    args = options.parse_args_and_arch(parser)
349
350
351

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

353
354
355
        distributed_main(args)
    elif args.distributed_world_size > 1:
        from multiprocessing_train import main as multiprocessing_main
356

357
358
359
        multiprocessing_main(args)
    else:
        main(args)