24 lines
619 B
Python
24 lines
619 B
Python
import torch
|
|
import ignite.distributed as idist
|
|
|
|
from omegaconf import OmegaConf
|
|
|
|
from model import MODEL
|
|
import torch.optim as optim
|
|
|
|
|
|
def build_model(cfg):
|
|
cfg = OmegaConf.to_container(cfg)
|
|
bn_to_sync_bn = cfg.pop("_bn_to_sync_bn", False)
|
|
model = MODEL.build_with(cfg)
|
|
if bn_to_sync_bn:
|
|
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
|
|
return idist.auto_model(model)
|
|
|
|
|
|
def build_optimizer(params, cfg):
|
|
assert "_type" in cfg
|
|
cfg = OmegaConf.to_container(cfg)
|
|
optimizer = getattr(optim, cfg.pop("_type"))(params=params, **cfg)
|
|
return idist.auto_optim(optimizer)
|