Build the learning rate scheduler.
(optimizer)
| 351 | |
| 352 | |
| 353 | def get_learning_rate_scheduler(optimizer): |
| 354 | """Build the learning rate scheduler.""" |
| 355 | args = get_args() |
| 356 | |
| 357 | # Iteration-based training. |
| 358 | if args.train_iters: |
| 359 | if args.lr_decay_iters is None: |
| 360 | args.lr_decay_iters = args.train_iters |
| 361 | decay_steps = args.lr_decay_iters * args.global_batch_size |
| 362 | if args.lr_warmup_fraction is not None: |
| 363 | warmup_steps = args.lr_warmup_fraction * decay_steps |
| 364 | else: |
| 365 | warmup_steps = args.lr_warmup_iters * args.global_batch_size |
| 366 | # Sample-based training. |
| 367 | elif args.train_samples: |
| 368 | # We need to set training iters for later use. Technically |
| 369 | # we need to adjust the training samples too (due to last |
| 370 | # batch being incomplete) but we leave it as is for now. |
| 371 | update_train_iters(args) |
| 372 | if args.lr_decay_samples is None: |
| 373 | args.lr_decay_samples = args.train_samples |
| 374 | decay_steps = args.lr_decay_samples |
| 375 | if args.lr_warmup_fraction is not None: |
| 376 | warmup_steps = args.lr_warmup_fraction * decay_steps |
| 377 | else: |
| 378 | warmup_steps = args.lr_warmup_samples |
| 379 | else: |
| 380 | raise Exception("either train-iters or train-samples should be provided.") |
| 381 | |
| 382 | lr_scheduler = AnnealingLR( |
| 383 | optimizer, |
| 384 | max_lr=args.lr, |
| 385 | min_lr=args.min_lr, |
| 386 | warmup_steps=warmup_steps, |
| 387 | decay_steps=decay_steps, |
| 388 | decay_style=args.lr_decay_style, |
| 389 | use_checkpoint_lr_scheduler=args.use_checkpoint_lr_scheduler, |
| 390 | override_lr_scheduler=args.override_lr_scheduler, |
| 391 | ) |
| 392 | |
| 393 | return lr_scheduler |
| 394 | |
| 395 | |
| 396 | def setup_model_and_optimizer(model_provider_func): |
no test coverage detected