steplr.py 315 B

1234567891011121314
  1. #! /usr/bin/python
  2. # -*- encoding: utf-8 -*-
  3. import torch
  4. def Scheduler(optimizer, test_interval, max_epoch, lr_decay, **kwargs):
  5. sche_fn = torch.optim.lr_scheduler.StepLR(optimizer, step_size=test_interval, gamma=lr_decay)
  6. lr_step = 'epoch'
  7. print('Initialised step LR scheduler')
  8. return sche_fn, lr_step