def create_model(configs):
    """Create model based on architecture name"""
    try:
        arch_parts = configs.arch.split('_')
        num_layers = int(arch_parts[-1])
    except:
        raise ValueError
    if 'fpn_resnet' in configs.arch:
        print('using ResNet architecture with feature pyramid')
        model = fpn_resnet.get_pose_net(
            num_layers=num_layers,
            heads=configs.heads,
            head_conv=configs.head_conv,
            imagenet_pretrained=configs.imagenet_pretrained)
    elif 'resnet' in configs.arch:
        print('using ResNet architecture')
        model = resnet.get_pose_net(
            num_layers=num_layers,
            heads=configs.heads,
            head_conv=configs.head_conv,
            imagenet_pretrained=configs.imagenet_pretrained)
    else:
        assert False, 'Undefined model backbone'

    return model
import sys
sys.path.insert(0, "/export/guanghan/CenterNet-Gluon/")
sys.path.insert(0, "/Users/guanghan.ning/Desktop/dev/CenterNet-Gluon/")

from models.model import create_model, load_model, save_model
from opts import opts

from models.resnet import get_pose_net
from mxnet import nd, gluon, init
import mxnet as mx

print('Creating model...')
opt = opts().init()
print(opt.arch)
model = get_pose_net(18, opt.heads, opt.head_conv, load_pretrained=True)

#model.collect_params().initialize(init=init.Xavier())

X = nd.random.uniform(shape=(16, 3, 512, 512))
print("\t Input shape: ", X.shape)
Y = model(X)
print("output: heatmaps", Y[0]["hm"].shape)
print("output: wh_scale", Y[0]["wh"].shape)
print("output: xy_offset", Y[0]["reg"].shape)
#print("output: xy_offset", Y[0]["reg"])

param = model.collect_params()
param_keys = param.keys()
print(param_keys)
#param_keys_residual_1 = [param[param_key] for param_key in param_keys if "hourglassnet0_residual1_conv1_weight" in param_key]
#print(param_keys_residual_1)