def get_net(cfg, weight_path=None):
"""
根据网络名称获取模型
:param net_name: 网络名称
:param weight_path: 与训练权重路径
:return:
"""
model = build_model(cfg)
if cfg.model_path is not None:
model.load_state_dict(torch.load(cfg.model_path))
return model
from inicls import build_model
from mmcv import Config
from tools.torch_utils import *
cfg = Config.fromfile('./config_test.py')
model = build_model(cfg=cfg)
model = model.cuda()
model.train()
print('[i] Architecture is {}'.format(cfg.model))
print('[i] Total Params: %.2fM' % (calculate_parameters(model)))