| import torch |
| from torch import optim |
|
|
|
|
| def build_optimizer(cfg, model, base_lr=0.0, resume=None): |
| print('==============================') |
| print('Optimizer: {}'.format(cfg['optimizer'])) |
| print('--momentum: {}'.format(cfg['momentum'])) |
| print('--weight_decay: {}'.format(cfg['weight_decay'])) |
|
|
| if cfg['optimizer'] == 'sgd': |
| optimizer = optim.SGD( |
| model.parameters(), |
| lr=base_lr, |
| momentum=cfg['momentum'], |
| weight_decay=cfg['weight_decay']) |
|
|
| elif cfg['optimizer'] == 'adam': |
| optimizer = optim.Adam( |
| model.parameters(), |
| lr=base_lr, |
| eight_decay=cfg['weight_decay']) |
| |
| elif cfg['optimizer'] == 'adamw': |
| optimizer = optim.AdamW( |
| model.parameters(), |
| lr=base_lr, |
| weight_decay=cfg['weight_decay']) |
| |
| start_epoch = 0 |
| if resume is not None: |
| print('keep training: ', resume) |
| checkpoint = torch.load(resume) |
| |
| checkpoint_state_dict = checkpoint.pop("optimizer") |
| optimizer.load_state_dict(checkpoint_state_dict) |
| start_epoch = checkpoint.pop("epoch") |
| |
| |
| return optimizer, start_epoch |
|
|
