Spaces:
Build error
Build error
File size: 1,501 Bytes
2c03f99 |
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 42 43 |
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
import os
import math
import torch
import torch.nn as nn
def build_optimizer(cfg, model):
""" Build optimizer from cfg file."""
g_bnw, g_w, g_b = [], [], []
for v in model.modules():
if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):
g_b.append(v.bias)
if isinstance(v, nn.BatchNorm2d):
g_bnw.append(v.weight)
elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):
g_w.append(v.weight)
assert cfg.solver.optim == 'SGD' or 'Adam', 'ERROR: unknown optimizer, use SGD defaulted'
if cfg.solver.optim == 'SGD':
optimizer = torch.optim.SGD(g_bnw, lr=cfg.solver.lr0, momentum=cfg.solver.momentum, nesterov=True)
elif cfg.solver.optim == 'Adam':
optimizer = torch.optim.Adam(g_bnw, lr=cfg.solver.lr0, betas=(cfg.solver.momentum, 0.999))
optimizer.add_param_group({'params': g_w, 'weight_decay': cfg.solver.weight_decay})
optimizer.add_param_group({'params': g_b})
del g_bnw, g_w, g_b
return optimizer
def build_lr_scheduler(cfg, optimizer, epochs):
"""Build learning rate scheduler from cfg file."""
if cfg.solver.lr_scheduler == 'Cosine':
lf = lambda x: ((1 - math.cos(x * math.pi / epochs)) / 2) * (cfg.solver.lrf - 1) + 1
else:
LOGGER.error('unknown lr scheduler, use Cosine defaulted')
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
return scheduler, lf
|