File size: 1,318 Bytes
beb7843 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
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_state_dict = checkpoint.pop("optimizer")
optimizer.load_state_dict(checkpoint_state_dict)
start_epoch = checkpoint.pop("epoch")
return optimizer, start_epoch
|