def prt_model(model_id, output_layer):
# ResNet50
print('loading model...')
# 0 ResNet50
if model_id == 0:
if output_layer == 0: # Keras
model_name = 'ResNet50_Keras'
model = ResNet50(weights='imagenet', include_top=False)
img_size = 224
#img_size = 30
model_path = '../models/resnet50.onnx'
datain = 'gpu_0/data_0'
if output_layer == 1: # ONNX - Last Layer
model_name = 'ResNet50_ONNX_L'
layerout = 'last'
elif output_layer == 2: # ONNX - Previous Layer
model_name = 'ResNet50_ONNX_P'
layerout = 'flatten0_output'
# 1 VGG16
elif model_id == 1:
if output_layer == 0: # Keras
model_name = 'VGG16_Keras'
model = VGG16(weights='imagenet', include_top=False)
#img_size = 224
img_size = 60
model_path = '../models/vgg16.onnx'
datain = 'data'
if output_layer == 1: # ONNX - Last Layer
model_name = 'VGG16_ONNX_L'
layerout = 'last'
elif output_layer == 2: # ONNX - Previous Layer
model_name = 'VGG16_ONNX_P'
layerout = 'flatten2_output'
# 2 VGG19
elif model_id == 2:
if output_layer == 0: # Keras
model_name = 'VGG19_Keras'
model = VGG19(weights='imagenet', include_top=False)
#img_size = 224
img_size = 60
model_path = '../models/vgg19.onnx'
datain = 'data_0'
if output_layer == 1: # ONNX - Last Layer
model_name = 'VGG19_ONNX_L'
layerout = 'last'
elif output_layer == 2: # ONNX - Previous Layer
model_name = 'VGG19_ONNX_P'
layerout = 'flatten2_output'
# 3 InceptionV3
elif model_id == 3:
if output_layer == 0: # Keras
model_name = 'InceptionV3_Keras'
model = InceptionV3(weights='imagenet', include_top=False)
#img_size = 224
img_size = 90
# 4 InceptionResNetV2
#elif model_id == 4:
# if output_layer == 0: # Keras
# model_name = 'InceptionResNetV2_Keras'
# model = InceptionResNetV2(weights='imagenet', include_top=False)
# #img_size = 224
# img_size = 90
# 5 Xception
elif model_id == 5:
if output_layer == 0: # Keras
model_name = 'Xception_Keras'
model = Xception(weights='imagenet', include_top=False)
#img_size = 224
img_size = 30
# 6 MobileNet
elif model_id == 6:
model_path = '../models/mobilenet.onnx'
datain = 'data'
if output_layer == 1: # ONNX - Last Layer
model_name = 'MobileNet_ONNX_L'
layerout = 'last'
elif output_layer == 2: # ONNX - Previous Layer
model_name = 'MobileNet_ONNX_P'
layerout = 'mobilenetv20_features_pool0_fwd_output' # not good
# 7 SqueezeNet
elif model_id == 7:
model_path = '../models/squeezenet.onnx'
datain = 'data'
if output_layer == 1: # ONNX - Last Layer
model_name = 'SqueezeNet_ONNX_L'
layerout = 'last'
elif output_layer == 2: # ONNX - Previous Layer
model_name = 'SqueezeNet_ONNX_P'
layerout = 'pooling3_output' # not good
# 8 AlexNet
elif model_id == 8:
model_path = '../models/alexnet.onnx'
datain = 'data_0'
if output_layer == 1: # ONNX - Last Layer
model_name = 'AlexNet_ONNX_L'
layerout = 'last'
elif output_layer == 2: # ONNX - Previous Layer
model_name = 'AlexNet_ONNX_P'
layerout = 'flatten2_output'
# 9 GoogleNet
elif model_id == 9:
model_path = '../models/googlenet.onnx'
datain = 'data_0'
if output_layer == 1: # ONNX - Last Layer
model_name = 'GoogleNet_ONNX_L'
layerout = 'last'
elif output_layer == 2: # ONNX - Previous Layer
model_name = 'GoogleNet_ONNX_P'
layerout = 'flatten0_output' # not good
# 10 ShuffleNet
elif model_id == 10:
model_path = '../models/shufflenet.onnx'
datain = 'gpu_0/data_0'
if output_layer == 1: # ONNX - Last Layer
model_name = 'ShuffleNet_ONNX_L'
layerout = 'last'
elif output_layer == 2: # ONNX - Previous Layer
model_name = 'ShuffleNet_ONNX_P'
layerout = 'flatten0_output' # not good
# 11 DenseNet121
elif model_id == 11:
model_path = '../models/densenet121.onnx'
datain = 'data_0'
if output_layer == 1: # ONNX - Last Layer
model_name = 'DenseNet121_ONNX_L'
layerout = 'last'
elif output_layer == 2: # ONNX - Previous Layer
model_name = 'DenseNet121_ONNX_P'
layerout = 'pad124_output' # not good
# 12 ZfNet512
elif model_id == 12:
model_path = '../models/zfnet512.onnx'
datain = 'gpu_0/data_0'
if output_layer == 1: # ONNX - Last Layer
model_name = 'ZfNet512_ONNX_L'
layerout = 'last'
elif output_layer == 2: # ONNX - Previous Layer
model_name = 'ZfNet512_ONNX_P'
layerout = 'flatten2_output' # not good
# 13 RCNN_ILSVRC13
elif model_id == 13:
model_path = '../models/rcnn_ilsvrc13.onnx'
datain = 'data_0'
if output_layer == 1: # ONNX - Last Layer
model_name = 'RCNN_ILSVRC13_ONNX_L'
layerout = 'last'
elif output_layer == 2: # ONNX - Previous Layer
model_name = 'RCNN_ILSVRC13_ONNX_P'
layerout = 'flatten2_output' # not good
if output_layer > 0:
print('loading model ' + model_path + '...')
sym, arg_params, aux_params = import_model(model_path)
if len(mx.test_utils.list_gpus()) == 0:
ctx = mx.cpu()
else:
ctx = mx.gpu(0)
all_layers = sym.get_internals()
print(all_layers.list_outputs())
if layerout == 'last':
sym3 = sym
else:
sym3 = all_layers[layerout]
model = mx.mod.Module(symbol=sym3,
context=ctx,
label_names=None,
data_names=[datain])
image_size = (224, 224)
img_size = image_size[1]
model.bind(data_shapes=[(datain, (1, 3, image_size[0],
image_size[1]))])
model.set_params(arg_params, aux_params)
print('model ' + model_name + ' loaded.')
return model, img_size, model_name
