INNER CODE UNIT · Python
num_params
nv-tlabs/GSCNN · network/__init__.py:13
num_params = sum([param.nelement() for param in net.parameters()])
logging.info('Model params = {:2.1f}M'.format(num_params / 1000000))
net = net.cuda()
net = torch.nn.DataParallel(net)
return net
def get_model(network, num_classes, criterion, trunk):
module = network[:network.rfind('.')]
model = network[network.rfind('.')+1:]
mod = importlib.import_module(module)
net_func = getattr(mod, model)
net = net_func(num_classes=num_classes, trunk=trunk, criterion=criterion)
return net