train.py 14 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
10
11
12
import collections
import os
import math
import torch
Sergey Edunov's avatar
Sergey Edunov committed
13

14
15
from itertools import islice

alexeib's avatar
alexeib committed
16
from fairseq import criterions, models, options, progress_bar
17
from fairseq.data import data_utils, data_loaders
18
19
20
from fairseq.fp16_trainer import FP16Trainer
from fairseq.trainer import Trainer
from fairseq.meters import AverageMeter, StopwatchMeter
21
from fairseq.utils import checkpoint_paths
Sergey Edunov's avatar
Sergey Edunov committed
22

Myle Ott's avatar
Myle Ott committed
23

Myle Ott's avatar
Myle Ott committed
24
def main(args):
25
26
27
    if args.max_tokens is None:
        args.max_tokens = 6000

28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
    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)

    # Load dataset
    splits = ['train', 'valid']
    dataset = load_dataset(args, splits)
    print('| [{}] dictionary: {} types'.format(dataset.src, len(dataset.src_dict)))
    print('| [{}] dictionary: {} types'.format(dataset.dst, len(dataset.dst_dict)))
    for split in splits:
        print('| {} {} {} examples'.format(args.data, split, len(dataset.splits[split])))

    model = models.build_model(args, dataset.src_dict, dataset.dst_dict)
44

45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
    criterion = criterions.build_criterion(args, dataset.src_dict, dataset.dst_dict)
    print('| model {}, criterion {}'.format(args.arch, criterion.__class__.__name__))
    print('| num. model params: {}'.format(sum(p.data.numel() for p in model.parameters())))

    # Build trainer
    if args.fp16:
        trainer = FP16Trainer(args, model, criterion)
    else:
        if torch.cuda.get_device_capability(0)[0] >= 7:
            print('| NOTICE: your device may support faster training with --fp16')
        trainer = Trainer(args, model, criterion)
    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
    train_dataloader = dataset.train_dataloader_generator(
        args.train_subset,
        max_tokens=args.max_tokens,
        max_sentences=args.max_sentences,
        max_positions=(
            min(args.max_source_positions, trainer.get_model().max_encoder_positions()),
            min(args.max_target_positions, trainer.get_model().max_decoder_positions())
        ),
        seed=args.seed,
        sample_without_replacement=args.sample_without_replacement,
        shard_id=args.distributed_rank,
        num_shards=args.distributed_world_size,
    )

    # Load the latest checkpoint if one is available
alexeib's avatar
alexeib committed
78
    epoch, next_ds = load_checkpoint(args, trainer, train_dataloader)
79
80

    # Send a dummy batch to warm the caching allocator
alexeib's avatar
alexeib committed
81
    dummy_batch = data_utils.get_dummy_batch(args.max_tokens, dataset.src_dict, dataset.dst_dict)
82
83
84
85
86
87
    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()
88
    first_val_loss = None
89
90
    train_meter = StopwatchMeter()
    train_meter.start()
91
    valid_subsets = args.valid_subset.split(',')
92
93
    while lr > args.min_lr and epoch <= max_epoch and trainer.get_num_updates() < max_update:
        # train for one epoch
alexeib's avatar
alexeib committed
94
        train(args, trainer, next_ds, epoch, dataset)
95
96

        if epoch % args.validate_interval == 0:
97
            valid_losses = validate(args, trainer, dataset, valid_subsets, epoch)
98
99

        # only use first validation loss to update the learning rate
100
        lr = trainer.lr_step(epoch, valid_losses[0])
101
102

        # save checkpoint
103
        if epoch % args.save_interval == 0:
104
            save_checkpoint(args, trainer, epoch, end_of_epoch=True, val_loss=valid_losses[0])
105
106

        epoch += 1
alexeib's avatar
alexeib committed
107
        next_ds = next(train_dataloader)
108
109
110
111
112
    train_meter.stop()
    print('| done training in {:.1f} seconds'.format(train_meter.sum))


def load_dataset(args, splits):
alexeib's avatar
alexeib committed
113
114
    is_raw = not data_utils.has_binary_files(args.data, splits)
    dataset = data_loaders.load_dataset(args, splits, is_raw)
115
116
117
    return dataset


118
def train(args, trainer, itr, epoch, dataset):
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
    """Train the model for one epoch."""

    # Set seed based on args.seed and the epoch number so that we get
    # reproducible results when resuming from checkpoints
    seed = args.seed + epoch
    torch.manual_seed(seed)

    # reset training meters
    for k in ['train_loss', 'train_nll_loss', 'wps', 'ups', 'wpb', 'bsz', 'clip']:
        meter = trainer.get_meter(k)
        if meter is not None:
            meter.reset()

    # update parameters every N batches
    if epoch <= len(args.update_freq):
        update_freq = args.update_freq[epoch - 1]
    else:
        update_freq = args.update_freq[-1]

    extra_meters = collections.defaultdict(lambda: AverageMeter())
139
    first_valid = args.valid_subset.split(',')[0]
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
    max_update = args.max_update or math.inf
    num_batches = len(itr)
    progress = progress_bar.build_progress_bar(args, itr, epoch, no_progress_bar='simple')
    for i, sample in enumerate(progress):
        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()

167
        num_updates = trainer.get_num_updates()
168
        if args.save_interval_updates > 0 and num_updates % args.save_interval_updates == 0:
169
170
            valid_losses = validate(args, trainer, dataset, [first_valid], epoch)
            save_checkpoint(args, trainer, epoch, end_of_epoch=False, val_loss=valid_losses[0])
171
172

        if num_updates >= max_update:
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
            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)


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)
    return stats


206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
def validate(args, trainer, dataset, subsets, epoch):
    """Evaluate the model on the validation set(s) and return the losses."""
    valid_losses = []
    for subset in subsets:
        # Initialize dataloader
        max_positions_valid = (
            trainer.get_model().max_encoder_positions(),
            trainer.get_model().max_decoder_positions(),
        )
        itr = dataset.eval_dataloader(
            subset,
            max_tokens=args.max_tokens,
            max_sentences=args.max_sentences_valid,
            max_positions=max_positions_valid,
            skip_invalid_size_inputs_valid_test=args.skip_invalid_size_inputs_valid_test,
            descending=True,  # largest batch first to warm the caching allocator
            shard_id=args.distributed_rank,
            num_shards=args.distributed_world_size,
        )
        progress = progress_bar.build_progress_bar(
            args, itr, epoch,
            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())
        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)
245

246
247
248
249
250
        # log validation stats
        stats = get_valid_stats(trainer)
        for k, meter in extra_meters.items():
            stats[k] = meter.avg
        progress.print(stats)
251

252
253
        valid_losses.append(stats['valid_loss'])
    return valid_losses
254
255
256
257
258
259
260
261


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
262
    else:
263
264
        nll_loss = trainer.get_meter('valid_loss').avg
    stats['valid_ppl'] = get_perplexity(nll_loss)
Myle Ott's avatar
Nits  
Myle Ott committed
265
266
267
    stats['num_updates'] = trainer.get_num_updates()
    if hasattr(save_checkpoint, 'best'):
        stats['best'] = min(save_checkpoint.best, stats['valid_loss'])
268
269
270
271
272
273
274
275
276
277
    return stats


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


278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
def save_checkpoint(args, trainer, epoch, end_of_epoch, val_loss):
    if args.no_save or args.distributed_rank > 0:
        return
    updates = trainer.get_num_updates()

    checkpoint_conds = collections.OrderedDict()
    checkpoint_conds['checkpoint{}.pt'.format(epoch)] = (
        end_of_epoch and not args.no_epoch_checkpoints and
        epoch % args.save_interval == 0
    )
    checkpoint_conds['checkpoint_{}_{}.pt'.format(epoch, updates)] = (
        not end_of_epoch and args.save_interval_updates > 0 and
        updates % args.save_interval_updates == 0
    )
    checkpoint_conds['checkpoint_best.pt'] = (
        not hasattr(save_checkpoint, 'best') or val_loss < save_checkpoint.best
    )
    checkpoint_conds['checkpoint_last.pt'] = True  # keep this last so that it's a symlink

    save_checkpoint.best = min(val_loss, getattr(save_checkpoint, 'best', val_loss))
298
    extra_state = {
299
300
        'best': save_checkpoint.best,
        'end_of_epoch': end_of_epoch,
301
302
        'epoch': epoch,
        'val_loss': val_loss,
303
        'wall_time': trainer.get_meter('wall').elapsed_time,
304
305
    }

306
307
308
309
310
311
312
313
314
315
316
317
    checkpoints = [os.path.join(args.save_dir, fn) for fn, cond in checkpoint_conds.items() if cond]
    if len(checkpoints) > 0:
        for fn in checkpoints:
            if os.path.exists(fn):
                os.remove(fn)
        trainer.save_checkpoint(checkpoints[0], extra_state)
        for fn in checkpoints[1:]:
            os.symlink(os.path.basename(checkpoints[0]), fn)

    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')
318
319
        for old_chk in checkpoints[args.keep_interval_updates:]:
            os.remove(old_chk)
320
321
322
323
324
325


def load_checkpoint(args, trainer, train_dataloader):
    os.makedirs(args.save_dir, exist_ok=True)
    checkpoint_path = os.path.join(args.save_dir, args.restore_file)
    epoch = 1
alexeib's avatar
alexeib committed
326
    ds = None
327
328
329
330
    if os.path.isfile(checkpoint_path):
        extra_state = trainer.load_checkpoint(checkpoint_path)
        if extra_state is not None:
            epoch = extra_state['epoch']
alexeib's avatar
alexeib committed
331
            end_of_epoch = extra_state.get('end_of_epoch', True)
alexeib's avatar
alexeib committed
332
333
            trainer_updates = trainer.get_num_updates()

334
335
336
            if 'best' in extra_state:
                save_checkpoint.best = extra_state['best']

alexeib's avatar
alexeib committed
337
            print('| loaded checkpoint {} (epoch {})'.format(checkpoint_path, epoch))
alexeib's avatar
alexeib committed
338

339
            trainer.lr_step(epoch)
alexeib's avatar
alexeib committed
340
            updates = 0
341
            for i in range(epoch):
alexeib's avatar
alexeib committed
342
343
344
                ds = next(train_dataloader)
                updates += len(ds)

alexeib's avatar
alexeib committed
345
            if not end_of_epoch and ds is not None and updates > trainer_updates:
346
347
348
349
                completed_batches = len(ds) - (updates - trainer_updates)
                assert completed_batches >= 0
                ds = iter(ds)

alexeib's avatar
alexeib committed
350
351
                print('| resuming from batch {}'.format(completed_batches + 1))

352
353
                # consume completed batches
                next(islice(ds, completed_batches, completed_batches), None)
alexeib's avatar
alexeib committed
354
            else:
alexeib's avatar
alexeib committed
355
356
                if not end_of_epoch:
                    print('| WARNING: checkpoint is not at end of epoch')
alexeib's avatar
alexeib committed
357
358
359
                ds = next(train_dataloader)
                epoch += 1

360
            trainer.get_meter('wall').reset(init=extra_state.get('wall_time', 0))
alexeib's avatar
alexeib committed
361
    return epoch, ds or next(train_dataloader)
Sergey Edunov's avatar
Sergey Edunov committed
362

Myle Ott's avatar
Myle Ott committed
363

Sergey Edunov's avatar
Sergey Edunov committed
364
if __name__ == '__main__':
Myle Ott's avatar
Myle Ott committed
365
366
    parser = options.get_training_parser()
    args = options.parse_args_and_arch(parser)
367
368
369

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

371
372
373
        distributed_main(args)
    elif args.distributed_world_size > 1:
        from multiprocessing_train import main as multiprocessing_main
374

375
376
377
        multiprocessing_main(args)
    else:
        main(args)