def test_inference():
all_models = ['resnet50_v1', 'vgg19_bn', 'alexnet', #'inceptionv3',
'densenet201', 'squeezenet1.0', 'mobilenet0.25']
batch_size = 10
download_data()
for model_name in all_models:
eprint('testing inference on %s'%model_name)
data_shape = (3, 224, 224) if 'inception' not in model_name else (3, 299, 299)
dataIter = mx.io.ImageRecordIter(
path_imgrec = VAL_DATA,
label_width = 1,
preprocess_threads = 1,
batch_size = batch_size,
data_shape = data_shape,
label_name = 'softmax_label',
rand_crop = False,
rand_mirror = False)
data_batch = dataIter.next()
data = data_batch.data[0]
label = data_batch.label[0]
gpu_data = data.as_in_context(mx.gpu())
gpu_label = label.as_in_context(mx.gpu())
# This is to create a model and run the model once to initialize
# all parameters.
cpu_model = get_model(model_name)
cpu_model.collect_params().initialize(ctx=mx.cpu())
cpu_model(mx.nd.array(data, ctx=mx.cpu()))
gpu_model = get_model(model_name)
gpu_model.collect_params().initialize(ctx=mx.gpu())
gpu_model(mx.nd.array(data, ctx=mx.gpu()))
# Force the two models have the same parameters.
cpu_params = cpu_model.collect_params()
gpu_params = gpu_model.collect_params()
for k in cpu_params.keys():
k = k.replace(cpu_params.prefix, '')
cpu_param = cpu_params.get(k)
gpu_param = gpu_params.get(k)
gpu_param.set_data(cpu_param.data().as_in_context(mx.gpu()))
for i in range(5):
# Run inference.
with autograd.record(train_mode=False):
cpu_out = cpu_model(mx.nd.array(data, ctx=mx.cpu()))
gpu_out = gpu_model(gpu_data)
out = cpu_out.asnumpy()
max_val = np.max(np.abs(out))
gpu_max_val = np.max(np.abs(gpu_out.asnumpy()))
eprint(model_name + ": CPU " + str(max_val) + ", GPU " + str(gpu_max_val))
assert_almost_equal(out / max_val, gpu_out.asnumpy() / max_val, rtol=1e-3, atol=1e-3)
def compile_network(opt, env, target):
# Populate the shape and data type dictionary
dtype_dict = {"data": "float32"}
shape_dict = {"data": (env.BATCH, 3, 224, 224)}
# Get off the shelf gluon model, and convert to relay
gluon_model = vision.get_model(opt.model, pretrained=True)
mod, params = relay.frontend.from_mxnet(gluon_model, shape_dict)
# Update shape and type dictionary
shape_dict.update({k: v.shape for k, v in params.items()})
dtype_dict.update({k: str(v.dtype) for k, v in params.items()})
# Perform quantization in Relay
# Note: We set opt_level to 3 in order to fold batch norm
with tvm.transform.PassContext(opt_level=3):
with relay.quantize.qconfig(global_scale=8.0, skip_conv_layers=[0]):
relay_prog = relay.quantize.quantize(mod["main"], params=params)
# Perform graph packing and constant folding for VTA target
if target.device_name == "vta":
assert env.BLOCK_IN == env.BLOCK_OUT
relay_prog = graph_pack(
relay_prog,
env.BATCH,
env.BLOCK_OUT,
env.WGT_WIDTH,
start_name=opt.start_name,
stop_name=opt.stop_name,
)
return relay_prog, params
def load_mxnet_resnet():
"""Load a pretrained resnet model from MXNet and transform that into NNVM
format.
Returns
-------
net : nnvm.Symbol
The loaded resnet computation graph.
params : dict[str -> NDArray]
The pretrained model parameters.
data_shape: tuple
The shape of the input tensor (an image).
out_shape: tuple
The shape of the output tensor (probability of all classes).
"""
print("Loading pretrained resnet model from MXNet...")
# Download a pre-trained mxnet resnet18_v1 model.
from mxnet.gluon.model_zoo.vision import get_model
block = get_model('resnet18_v1', pretrained=True)
# Transform the mxnet model into NNVM.
# We want a probability so add a softmax operator.
sym, params = nnvm.frontend.from_mxnet(block)
sym = nnvm.sym.softmax(sym)
print("- Model loaded!")
return sym, params, (1, 3, 224, 224), (1, 1000)
def get_network(name, batch_size):
"""Get the symbol definition and random weight of a network"""
input_shape = (batch_size, 3, 224, 224)
output_shape = (batch_size, 1000)
if "resnet" in name:
n_layer = int(name.split('-')[1])
net, params = relay.testing.resnet.get_workload(num_layers=n_layer, batch_size=batch_size, dtype=dtype)
elif "vgg" in name:
n_layer = int(name.split('-')[1])
net, params = relay.testing.vgg.get_workload(num_layers=n_layer, batch_size=batch_size, dtype=dtype)
elif name == 'mobilenet':
net, params = relay.testing.mobilenet.get_workload(batch_size=batch_size, dtype=dtype)
elif name == 'squeezenet_v1.1':
net, params = relay.testing.squeezenet.get_workload(batch_size=batch_size, version='1.1', dtype=dtype)
elif name == 'inception_v3':
input_shape = (1, 3, 299, 299)
net, params = relay.testing.inception_v3.get_workload(batch_size=batch_size, dtype=dtype)
elif name == 'mxnet':
# an example for mxnet model
from mxnet.gluon.model_zoo.vision import get_model
block = get_model('resnet18_v1', pretrained=True)
net, params = relay.frontend.from_mxnet(block, shape={'data': input_shape}, dtype=dtype)
net = relay.Function(net.params, relay.nn.softmax(net.body), None, net.type_params, net.attrs)
else:
raise ValueError("Unsupported network: " + name)
return net, params, input_shape, output_shape
def get_model(model_name, batch_size=1):
if model_name == "resnet18_v1":
import mxnet as mx
from mxnet import gluon
from mxnet.gluon.model_zoo import vision
gluon_model = vision.get_model(model_name, pretrained=True)
img_size = 224
data_shape = (batch_size, 3, img_size, img_size)
net, params = relay.frontend.from_mxnet(gluon_model, {"data": data_shape})
return (net, params)
elif model_name == "mobilenet_v2":
import keras
from keras.applications.mobilenet_v2 import MobileNetV2
keras.backend.clear_session() # Destroys the current TF graph and creates a new one.
weights_url = "".join(
[
"https://github.com/JonathanCMitchell/",
"mobilenet_v2_keras/releases/download/v1.1/",
"mobilenet_v2_weights_tf_dim_ordering_tf_kernels_0.5_224.h5",
]
)
weights_file = "mobilenet_v2_weights.h5"
weights_path = download_testdata(weights_url, weights_file, module="keras")
keras_mobilenet_v2 = MobileNetV2(
alpha=0.5, include_top=True, weights=None, input_shape=(224, 224, 3), classes=1000
)
keras_mobilenet_v2.load_weights(weights_path)
img_size = 224
data_shape = (batch_size, 3, img_size, img_size)
mod, params = relay.frontend.from_keras(keras_mobilenet_v2, {"input_1": data_shape})
return (mod, params)
def get_net(model_name, context):
if model_name == 'FCN':
num_classes = 21
pretrained_net = vision.resnet18_v2(pretrained=True, ctx=context)
net = gluon.nn.HybridSequential()
for layer in pretrained_net.features[:-2]:
net.add(layer)
net.add(
gluon.nn.Conv2D(num_classes, kernel_size=1),
gluon.nn.Conv2DTranspose(num_classes,
kernel_size=64,
padding=16,
strides=32))
net[-1].initialize(init.Constant(
bilinear_kernel(num_classes, num_classes, 64)),
ctx=context)
net[-2].initialize(init=init.Xavier(), ctx=context)
input_shape = (1, 3, 320, 480)
label_shape = (1, 320, 480)
loss_axis = 1
else:
net = vision.get_model(model_name, classes=10)
net.initialize(mx.init.Xavier(), ctx=context)
input_shape = (1, 1, 224, 224)
label_shape = 1
loss_axis = -1
return net, input_shape, label_shape, loss_axis
def train(network, batch_size, ctx):
assert (batch_size >= len(ctx)), "ERROR: batch size should not be smaller than num of GPUs."
net = models.get_model(network)
if 'inceptionv3' == network:
data_shape = [('data', (batch_size,) + image_shapes[1])]
else:
data_shape = [('data', (batch_size,) + image_shapes[0])]
data = mx.sym.var('data')
out = net(data)
softmax = mx.sym.SoftmaxOutput(out, name='softmax')
mod = mx.mod.Module(softmax, context=ctx)
mod.bind(for_training = True,
inputs_need_grad = False,
data_shapes = data_shape)
mod.init_params(initializer=mx.init.Xavier(magnitude=2.))
if len(ctx) > 1:
mod.init_optimizer(kvstore='device', optimizer='sgd')
else:
mod.init_optimizer(kvstore='local', optimizer='sgd')
data = [mx.random.uniform(-1.0, 1.0, shape=shape, ctx=ctx[0]) for _, shape in mod.data_shapes]
batch = mx.io.DataBatch(data, [])
for i in range(dry_run + num_batches):
if i == dry_run:
tic = time.time()
mod.forward(batch, is_train=True)
for output in mod.get_outputs():
output.wait_to_read()
mod.backward()
mod.update()
bwd = time.time() - tic
return bwd
def get_network(name, batch_size):
"""Get the symbol definition and random weight of a network"""
input_shape = (batch_size, 3, 224, 224)
output_shape = (batch_size, 1000)
if name =='resnet-18':
net, params = nnvm.testing.resnet.get_workload(num_layers=18, batch_size=batch_size)
elif name =='mobilenet':
net, params = nnvm.testing.mobilenet.get_workload(batch_size=batch_size)
elif name =='squeezenet v1.1':
net, params = nnvm.testing.squeezenet.get_workload(batch_size=batch_size, version='1.1')
elif name =='vgg-16':
net, params = nnvm.testing.vgg.get_workload(num_layers=16, batch_size=batch_size)
elif name =='custom':
# an example for custom network
from nnvm.testing import utils
net = nnvm.sym.Variable('data')
net = nnvm.sym.conv2d(net, channels=4, kernel_size=(3,3), padding=(1,1))
net = nnvm.sym.flatten(net)
net = nnvm.sym.dense(net, units=1000)
net, params = utils.create_workload(net, batch_size, (3, 224, 224))
elif name == 'mxnet':
# an example for mxnet model
from mxnet.gluon.model_zoo.vision import get_model
block = get_model('resnet18_v1', pretrained=True)
net, params = nnvm.frontend.from_mxnet(block)
net = nnvm.sym.softmax(net)
else:
raise ValueError("Unsupported network: " + name)
return net, params, input_shape, output_shape
def mobilenet():
with tvm.autotvm.tophub.context("llvm"):
dtype_dict = {"data": 'float32'}
shape_dict = {"data": (1, 3, 224, 224)}
model = vision.get_model("mobilenet0.25", pretrained=True)
mod, params = relay.frontend.from_mxnet(model, shape_dict)
print(mod)
return mod, params
def get_gluon_output(name, x):
net = vision.get_model(name)
net.collect_params().initialize(mx.init.Xavier())
net_sym = gluon.nn.SymbolBlock(outputs=net(mx.sym.var('data')),
inputs=mx.sym.var('data'),
params=net.collect_params())
out = net_sym(mx.nd.array(x.astype(dtype))).asnumpy()
return out, net_sym
def get_from_gluon(model_name, classes=1000, logger=None):
dir_path = os.path.dirname(os.path.realpath(__file__))
model_path = os.path.join(dir_path, 'model')
if logger is not None:
logger.info('Converting model from Gluon-CV ModelZoo %s... into path %s' % (model_name, model_path))
net = get_model(name=model_name, classes=classes, pretrained=True)
prefix = os.path.join(model_path, model_name)
return net, prefix
def get_gluon_output(name, x):
net = vision.get_model(name)
net.collect_params().initialize(mx.init.Xavier())
net_sym = gluon.nn.SymbolBlock(outputs=net(mx.sym.var('data')),
inputs=mx.sym.var('data'),
params=net.collect_params())
out = net_sym(mx.nd.array(x.astype(dtype))).asnumpy()
return out, net_sym
def test_tensorrt_serialize_graph_runtime():
if skip_codegen_test():
return
import mxnet as mx
from mxnet.gluon.model_zoo.vision import get_model
data_shape = (1, 3, 224, 224)
data_type = "float32"
i_data = np.random.uniform(0, 1, data_shape).astype(data_type)
block = get_model("resnet18_v1", pretrained=True)
mod, params = relay.frontend.from_mxnet(block,
shape={"data": data_shape},
dtype=data_type)
mod, config = tensorrt.partition_for_tensorrt(mod)
tmpdir = utils.tempdir()
def compile_graph(mod, params):
with tvm.transform.PassContext(
opt_level=3, config={"relay.ext.tensorrt.options": config}):
graph, lib, params = relay.build(mod, params=params, target="cuda")
params = relay.save_param_dict(params)
return graph, lib, params
def run_graph(graph, lib, params):
mod_ = graph_runtime.create(graph, lib, ctx=tvm.gpu(0))
mod_.load_params(params)
mod_.run(data=i_data)
res = mod_.get_output(0)
return res
def save_graph(graph, lib, params):
# Serialize
with open(tmpdir.relpath("compiled.json"), "w") as f_graph_json:
f_graph_json.write(graph)
with open(tmpdir.relpath("compiled.params"), "wb") as f_params:
f_params.write(params)
lib.export_library(tmpdir.relpath("compiled.so"))
def load_graph():
# Deserialize
with open(tmpdir.relpath("compiled.json"), "r") as f_graph_json:
graph = f_graph_json.read()
with open(tmpdir.relpath("compiled.params"), "rb") as f_params:
params = bytearray(f_params.read())
lib = tvm.runtime.load_module(tmpdir.relpath("compiled.so"))
return graph, lib, params
# Test serialization with graph runtime
graph, lib, graph_params = compile_graph(mod, params)
save_graph(graph, lib, graph_params)
loaded_graph, loaded_lib, loaded_params = load_graph()
if not skip_runtime_test():
result_dict = dict()
result_dict["graph"] = run_graph(graph, lib, graph_params)
result_dict["graph_ref"] = run_graph(loaded_graph, loaded_lib,
loaded_params)
assert_result_dict_holds(result_dict)
def run_and_verify_model(model):
if skip_codegen_test():
return
import mxnet as mx
from mxnet.gluon.model_zoo.vision import get_model
def check_trt_used(mod):
num_trt_subgraphs = sum([
1 if gv.name_hint == "tensorrt_0" else 0
for gv in mod.get_global_vars()
])
assert num_trt_subgraphs == 1
def compile_and_run(mod, params, i_data, mode="vm", use_trt=True):
assert mode in ["graph", "vm"]
if use_trt:
mod, config = tensorrt.partition_for_tensorrt(mod, params)
check_trt_used(mod)
with tvm.transform.PassContext(
opt_level=3, config={"relay.ext.tensorrt.options":
config}):
exec = relay.create_executor(mode,
mod=mod,
ctx=tvm.gpu(0),
target="cuda")
else:
with tvm.transform.PassContext(opt_level=3):
exec = relay.create_executor(mode,
mod=mod,
ctx=tvm.gpu(0),
target="cuda")
res = exec.evaluate()(i_data, **
params) if not skip_runtime_test() else None
return res
dtype = "float32"
input_shape = (1, 3, 224, 224)
i_data = np.random.uniform(-1, 1, input_shape).astype(dtype)
block = get_model(model, pretrained=True)
mod, params = relay.frontend.from_mxnet(block,
shape={"data": input_shape},
dtype=dtype)
result_dict = dict()
for mode in ["vm", "graph"]:
for use_trt in [True, False]:
result_key = mode + ("_trt" if use_trt else "")
result_dict[result_key] = compile_and_run(mod,
params,
i_data,
mode=mode,
use_trt=use_trt)
if not skip_runtime_test():
assert_result_dict_holds(result_dict)
def get_net(gpu_id):
param_path = 'checkpoints/Fashion_In.params'
base_net = vision_model.get_model('resnet50_v2')
net = MarginNet(base_net.features, 128, batch_k=5)
context = [mxnet.gpu(gpu_id)]
net.initialize()
net.collect_params().reset_ctx(context)
net.load_parameters(filename=param_path, ctx=context[0])
return net, context
def test_resnet_mxnet(env):
with tvm.autotvm.tophub.context("llvm"):
dtype_dict = {"data": 'float32'}
shape_dict = {"data": (env.BATCH, 3, 224, 224)}
model = vision.get_model("resnet18_v1", pretrained=True)
mod, params = relay.frontend.from_mxnet(model, shape_dict)
return mod, params
def build_model(model_name, remote, target, ctx, vta_env):
"""Build the inference graph runtime."""
# Load pre-configured AutoTVM schedules.
with autotvm.tophub.context(target):
# Populate the shape and data type dictionary for ResNet input.
dtype_dict = {'data': 'float32'}
shape_dict = {'data': (vta_env.BATCH, 3, 224, 224)}
# Get off-the-shelf gluon model and convert to Relay.
gluon_model = vision.get_model(model_name, pretrained=True)
# Start frontend compilation.
mod, params = relay.frontend.from_mxnet(gluon_model, shape_dict)
# Update shape and type dictionary.
shape_dict.update({k: v.shape for k, v in params.items()})
dtype_dict.update({k: str(v.dtype) for k, v in params.items()})
# Perform quantization in Relay.
with relay.quantize.qconfig(global_scale=8.0, skip_conv_layers=[0]):
relay_prog = relay.quantize.quantize(mod['main'], params=params)
# Perform graph packing and constant folding for VTA target.
if target.device_name == 'vta':
assert vta_env.BLOCK_IN == vta_env.BLOCK_OUT
relay_prog = graph_pack(relay_prog,
vta_env.BATCH,
vta_env.BLOCK_OUT,
vta_env.WGT_WIDTH,
start_name=START_PACK,
stop_name=STOP_PACK)
# Compile Relay program with AlterOpLayout disabled.
with relay.build_config(opt_level=3, disabled_pass={'AlterOpLayout'}):
if target.device_name == 'vta':
with vta.build_config():
graph, lib, params = relay.build(
relay_prog,
target=vta_env.target,
params=params,
target_host=vta_env.target_host)
else:
graph, lib, params = relay.build(
relay_prog,
target=target,
params=params,
target_host=vta_env.target_host)
# Send the inference library over to the remote RPC server
temp = util.tempdir()
lib.save(temp.relpath('graphlib.o'))
remote.upload(temp.relpath('graphlib.o'))
lib = remote.load_module('graphlib.o')
graph_module = graph_runtime.create(graph, lib, ctx)
graph_module.set_input(**params)
return graph_module
def resnet50_faster_rcnn(special_pretrain=False, load_path=None):
if special_pretrain and load_path is not None:
resnet = get_model('resnet50_v2', pretrained=False)
resnet.load_parameters(load_path)
else:
resnet = get_model('resnet50_v2', pretrained=True)
# with open('class_name', 'rb') as f:
# classes = pickle.load(f)
classes = ('dog', )
features = ResNet50_v2_features(resnet)
top_features_ = top_features()
model = FasterRCNN(features=features,
stride=32,
top_features=top_features_,
classes=classes,
short=600,
max_size=1000)
return model
def _parse_network(network, outputs, inputs, pretrained, ctx):
"""Parse network with specified outputs and other arguments.
Parameters
----------
network : str or HybridBlock or Symbol
Logic chain: load from gluon.model_zoo.vision if network is string.
Convert to Symbol if network is HybridBlock
outputs : str or iterable of str
The name of layers to be extracted as features.
inputs : iterable of str
The name of input datas.
pretrained : bool
Use pretrained parameters as in gluon.model_zoo
ctx : Context
The context, e.g. mxnet.cpu(), mxnet.gpu(0).
Returns
-------
inputs : list of Symbol
Network input Symbols, usually ['data']
outputs : list of Symbol
Network output Symbols, usually as features
params : ParameterDict
Network parameters.
"""
inputs = list(inputs) if isinstance(inputs, tuple) else inputs
for i, inp in enumerate(inputs):
if isinstance(inp, string_types):
inputs[i] = mx.sym.var(inp)
assert isinstance(inputs[i],
Symbol), "Network expects inputs are Symbols."
if len(inputs) == 1:
inputs = inputs[0]
else:
inputs = mx.sym.Group(inputs)
params = None
prefix = ''
if isinstance(network, string_types):
network = vision.get_model(network, pretrained=pretrained, ctx=ctx)
if isinstance(network, HybridBlock):
params = network.collect_params()
prefix = network._prefix
network = network(inputs)
assert isinstance(network, Symbol), \
"FeatureExtractor requires the network argument to be either " \
"str, HybridBlock or Symbol, but got %s"%type(network)
if isinstance(outputs, string_types):
outputs = [outputs]
assert len(outputs) > 0, "At least one outputs must be specified."
outputs = [
out if out.endswith('_output') else out + '_output' for out in outputs
]
outputs = [network.get_internals()[prefix + out] for out in outputs]
return inputs, outputs, params
def make_subgraph4(stype):
model = get_model('resnet18_v1')
model.hybridize()
model.initialize()
s = (1, 3, 32, 32)
data = mx.nd.random.normal(shape=s)
out = model(data)
model.export('resnet18')
orig = mx.sym.load('resnet18-symbol.json')
return make_subgraph_weight(orig, s, stype)
def get_tvm_workload(network, **kwargs):
from nnvm.frontend import from_mxnet
from mxnet.gluon.model_zoo.vision import get_model
block = get_model(network, **kwargs)
if network.startswith('resnet152'):
import sys
sys.setrecursionlimit(10000)
sym, params = from_mxnet(block)
sym = nnvm.sym.softmax(sym)
return sym, params
def download_block_from_mxnet(args):
import mxnet as mx
block = get_model(args.model, pretrained=True)
sym, params = nnvm.frontend.from_mxnet(block)
target = 'llvm'
shape = (1, 3, 224, 224) #input shape; need to be obtained from the model
shape_dict = {'data': shape}
model_layers = get_dataflow(sym, target, shape_dict, params, args.batch_size)
return model_layers
def model_fn(model_dir):
"""
Load the gluon model. Called once when hosting service starts.
:param: model_dir The directory where model files are stored.
:return: a model (in this case a Gluon network)
"""
net = models.get_model('vgg16_bn', ctx=mx.cpu(), pretrained=False, classes=10)
net.load_params('%s/model.params' % model_dir, ctx=mx.cpu())
return net
def creat_model(params_path, num_classes, anchor_scales):
pretrained = vision.get_model('resnet18_v1').features
net = nn.HybridSequential()
for i in range(len(pretrained) - 2):
net.add(pretrained[i])
predictor = Yolo2_Output(num_classes=num_classes,
anchor_scales=anchor_scales)
net.add(predictor)
net.load_parameters(params_path)
return net
def make_subgraph4(stype):
model = get_model('resnet18_v1')
model.hybridize()
model.initialize()
s = (1, 3, 32, 32)
data = mx.nd.random.normal(shape=s)
out = model(data)
model.export('resnet18')
orig = mx.sym.load('resnet18-symbol.json')
return make_subgraph_weight(orig, s, stype)
def get_network(name, batch_size, dtype='float32'):
"""Get the symbol definition and random weight of a network
Parameters
----------
name: str
The name of the network, can be 'resnet-18', 'resnet-50', 'vgg-16', 'inception_v3', 'mobilenet', ...
batch_size: int
batch size
dtype: str
Data type
Returns
-------
net: tvm.IRModule
The relay function of network definition
params: dict
The random parameters for benchmark
input_shape: tuple
The shape of input tensor
output_shape: tuple
The shape of output tensor
"""
input_shape = (batch_size, 3, 224, 224)
output_shape = (batch_size, 1000)
if name == 'mobilenet':
net, params = testing.mobilenet.get_workload(batch_size=batch_size, dtype=dtype)
elif name == 'inception_v3':
input_shape = (batch_size, 3, 299, 299)
net, params = testing.inception_v3.get_workload(batch_size=batch_size, dtype=dtype)
elif "resnet" in name:
n_layer = int(name.split('-')[1])
net, params = testing.resnet.get_workload(num_layers=n_layer, batch_size=batch_size, dtype=dtype)
elif "vgg" in name:
n_layer = int(name.split('-')[1])
net, params = testing.vgg.get_workload(num_layers=n_layer, batch_size=batch_size, dtype=dtype)
elif "densenet" in name:
n_layer = int(name.split('-')[1])
net, params = testing.densenet.get_workload(densenet_size=n_layer, batch_size=batch_size, dtype=dtype)
elif "squeezenet" in name:
version = name.split("_v")[1]
net, params = testing.squeezenet.get_workload(batch_size=batch_size, version=version, dtype=dtype)
elif name == 'mxnet':
# an example for mxnet model
from mxnet.gluon.model_zoo.vision import get_model
block = get_model('resnet18_v1', pretrained=True)
net, params = relay.frontend.from_mxnet(block, shape={'data': input_shape}, dtype=dtype)
net = net["main"]
net = relay.Function(net.params, relay.nn.softmax(net.body), None, net.type_params, net.attrs)
net = tvm.IRModule.from_expr(net)
else:
raise ValueError("Unsupported network: " + name)
return net, params, input_shape, output_shape
def get_network(name, batch_size):
"""Get the symbol definition and random weight of a network"""
input_shape = (batch_size, 3, 224, 224)
output_shape = (batch_size, 1000)
if "resnet" in name:
n_layer = int(name.split("-")[1])
mod, params = relay.testing.resnet.get_workload(
num_layers=n_layer, batch_size=batch_size, dtype=dtype
)
print("Tuning ResNet")
elif "vgg" in name:
n_layer = int(name.split("-")[1])
mod, params = relay.testing.vgg.get_workload(
num_layers=n_layer, batch_size=batch_size, dtype=dtype
)
print("Tuning VGG")
elif name == "mobilenet":
mod, params = relay.testing.mobilenet.get_workload(batch_size=batch_size, dtype=dtype)
print("Tuning MobileNet")
elif name == "squeezenet_v1.1":
mod, params = relay.testing.squeezenet.get_workload(
batch_size=batch_size, version="1.1", dtype=dtype
)
print("Tuning SqueezeNet")
elif name == "inception_v3":
input_shape = (1, 3, 299, 299)
mod, params = relay.testing.inception_v3.get_workload(batch_size=batch_size, dtype=dtype)
print("Tuning Inception")
elif name == "mxnet":
# an example for mxnet model
from mxnet.gluon.model_zoo.vision import get_model
if submod == 0:
modn = "resnet18_v1"
print("Tuning MXNet's ResNet")
elif submod == 1:
modn = "inceptionv3"
print("Tuning MXNet's Inception")
elif submod == 2:
modn = "mobilenetv2_1.0"
print("Tuning MXNet's MobileNet")
else:
raise Exception("Not Supported!")
block = get_model(modn, pretrained=True)
mod, params = relay.frontend.from_mxnet(block, shape={input_name: input_shape}, dtype=dtype)
net = mod["main"]
net = relay.Function(
net.params, relay.nn.softmax(net.body), None, net.type_params, net.attrs
)
mod = tvm.IRModule.from_expr(net)
else:
raise ValueError("Unsupported network: " + name)
return mod, params, input_shape, output_shape
def model_fn(model_dir):
"""
Load the gluon model. Called once when hosting service starts.
:param: model_dir The directory where model files are stored.
:return: a model (in this case a Gluon network)
"""
net = models.get_model('resnet34_v2', ctx=mx.cpu(), pretrained=False, classes=10)
net.load_params('%s/model.params' % model_dir, ctx=mx.cpu())
return net
def get_base_net(self):
"""
获取base net，默认是mobilenet v2
:return: base net，基础网络
"""
base_net = get_model('mobilenet1.0', pretrained=True)
with base_net.name_scope():
base_net.output = Dense(units=self.n_class) # 全连接层
base_net.output.initialize(Xavier(), ctx=self.ctx) # 初始化
base_net.collect_params().reset_ctx(self.ctx)
base_net.hybridize()
return base_net
def test_models(model_name):
pretrained_to_test = set(['mobilenetv2_0.25'])
test_pretrain = model_name in pretrained_to_test
model = get_model(model_name, pretrained=test_pretrain, root='model/')
data_shape = (2, 3, 224,
224) if 'inception' not in model_name else (2, 3, 299, 299)
eprint('testing forward for %s' % model_name)
print(model)
if not test_pretrain:
model.initialize()
model(mx.nd.random.uniform(shape=data_shape)).wait_to_read()
def end2end_benchmark(model, target, batch_size):
num_classes = 1000
image_shape = (3, 224, 224)
data_shape = (batch_size, ) + image_shape
out_shape = (batch_size, num_classes)
block = get_model(model, pretrained=True)
data_array = np.random.uniform(0, 255, size=data_shape).astype("float32")
times = []
for i in range(1):
s = time.time()
mxnet_out = block(mx.nd.array(data_array))
mxnet_out.asnumpy()
mkl_time = time.time() - s
times.append(mkl_time)
print("MKL %s inference time for batch size of %d: %f" %
(model, batch_size, np.mean(times) * 1000))
# data_array_5D = np.expand_dims(data_array, 4)
# print("data_array_5D shape: " + str(data_array_5D.shape))
net, params = nnvm.frontend.from_mxnet(block)
ctx = tvm.cpu()
opt_level = 2
with nnvm.compiler.build_config(opt_level=opt_level):
graph, lib, params = nnvm.compiler.build(net,
target,
shape={"data": data_shape},
params=params)
with open('graph.json', 'w') as fn:
fn.writelines(graph.json())
module = graph_runtime.create(graph, lib, ctx)
module.set_input(**params)
input_data = tvm.nd.array(data_array, ctx=ctx)
module.set_input('data', input_data)
times = []
for i in range(num_pass):
s = time.time()
module.run()
tvm_time = time.time() - s
times.append(tvm_time)
print("TVM %s inference time for batch size of %d: %f" %
(model, batch_size, np.mean(times) * 1000))
tvm_out = module.get_output(0, out=tvm.nd.empty(out_shape))
# decimal=3 does not work for resnet-101
np.testing.assert_array_almost_equal(tvm_out.asnumpy(),
mxnet_out.asnumpy(),
decimal=2)
return tvm_time, mkl_time
def create_model(self):
"""
Create the model
:return:
"""
tr_db = list(self.cfg['train'].values())[0]
num_class = tr_db['num_cls']
self.inference = get_model(self.args.arch, classes=num_class)
if self.args.hybridize:
self.inference.hybridize()
def get_network(name, batch_size):
"""Get the symbol definition and random weight of a network"""
input_shape = (batch_size, 3, 299, 299)
output_shape = (batch_size, 6)
if "resnet" in name:
n_layer = int(name.split('-')[1])
mod, params = relay.testing.resnet.get_workload(num_layers=n_layer,
batch_size=batch_size,
dtype=dtype)
elif "vgg" in name:
n_layer = int(name.split('-')[1])
mod, params = relay.testing.vgg.get_workload(num_layers=n_layer,
batch_size=batch_size,
dtype=dtype)
elif name == 'mobilenet':
mod, params = relay.testing.mobilenet.get_workload(
batch_size=batch_size, dtype=dtype)
elif name == 'squeezenet_v1.1':
mod, params = relay.testing.squeezenet.get_workload(
batch_size=batch_size, version='1.1', dtype=dtype)
elif name == 'inception_v3':
input_shape = (1, 3, 299, 299)
mod, params = relay.testing.inception_v3.get_workload(
batch_size=batch_size, dtype=dtype)
elif name == 'mxnet':
# an example for mxnet model
from mxnet.gluon.model_zoo.vision import get_model
block = get_model('resnet18_v1', pretrained=True)
mod, params = relay.frontend.from_mxnet(
block, shape={input_name: input_shape}, dtype=dtype)
net = mod["main"]
net = relay.Function(net.params, relay.nn.softmax(net.body), None,
net.type_params, net.attrs)
mod = tvm.IRModule.from_expr(net)
elif name == 'xception':
# an example for keras model (local model)
input_shape = (1, 3, 299, 299)
LOAD_JSON_PATH = '/Users/oldpan/Downloads/Xception_v0_2.json'
LOAD_MODEL_PATH = '/Users/oldpan/Downloads/Xception_v0_2.h5'
model = model_from_json(open(LOAD_JSON_PATH, 'r').read())
model.load_weights(LOAD_MODEL_PATH)
mod, params = relay.frontend.from_keras(
model, shape={input_name: input_shape})
net = mod["main"]
net = relay.Function(net.params, relay.nn.softmax(net.body), None,
net.type_params, net.attrs)
mod = tvm.IRModule.from_expr(net)
else:
raise ValueError("Unsupported network: " + name)
return mod, params, input_shape, output_shape
def get_base_net(ctx):
"""
获取base net，默认是mobilenet v2
:param ctx: 运行环境，cpu or gpu，默认cpu
:return: base net，基础网络
"""
base_net = get_model('mobilenet1.0', pretrained=True)
with base_net.name_scope():
base_net.output = Dense(units=27) # 全连接层
base_net.output.initialize(Xavier(), ctx=ctx) # 初始化
base_net.collect_params().reset_ctx(ctx)
base_net.hybridize()
return base_net
def load_model(self,weightFile = None):
bodyNet = vision.get_model('resnet18_v1', pretrained=True).features
net = nn.HybridSequential()
for k in range(len(bodyNet)-2):
net.add(bodyNet[k])
output = YOLO_OUTPUT(self.numClass, self.box_per_cell,verbose=False)
output.initialize()
net.add(output)
net.collect_params().reset_ctx(self.ctx)
if weightFile is not None:
net.load_params(weightFile,ctx=self.ctx)
self.net = net
return True
def compile_and_run(i_data,
input_shape,
dtype,
use_trt=True,
num_iteration=1):
import mxnet as mx
from mxnet.gluon.model_zoo.vision import get_model
def check_trt_used(graph):
import json
graph = json.loads(graph)
num_trt_subgraphs = sum([
1 for n in graph["nodes"] if n.get("attrs", {}).get(
"func_name", "").startswith("tensorrt_")
])
assert num_trt_subgraphs >= 1
block = get_model(model, pretrained=True)
mod, params = relay.frontend.from_mxnet(block,
shape={"data": input_shape},
dtype=dtype)
if use_trt:
mod, config = tensorrt.partition_for_tensorrt(mod, params)
with tvm.transform.PassContext(
opt_level=3, config={"relay.ext.tensorrt.options":
config}):
graph, lib, params = relay.build(mod, "cuda", params=params)
check_trt_used(graph)
else:
with tvm.transform.PassContext(opt_level=3):
graph, lib, params = relay.build(mod, "cuda", params=params)
if skip_runtime_test():
return
mod = graph_runtime.create(graph, lib, ctx=tvm.gpu(0))
mod.set_input(**params)
# Warmup
for i in range(10):
mod.run(data=i_data)
# Time
times = []
for i in range(num_iteration):
start_time = time.time()
mod.run(data=i_data)
res = mod.get_output(0)
times.append(time.time() - start_time)
latency = 1000.0 * np.mean(times)
print(model, latency)
return res
def test_model(model,
i_data,
input_shape,
dtype,
use_trt=True,
num_iteration=1000):
import mxnet
from mxnet.gluon.model_zoo.vision import get_model
def check_trt_used(graph):
import json
graph = json.loads(graph)
num_trt_subgraphs = sum([
1 for n in graph['nodes']
if n.get('attrs', {}).get('func_name', '') == 'tensorrt_0'
])
assert num_trt_subgraphs == 1
block = get_model(model, pretrained=True)
mod, params = relay.frontend.from_mxnet(block,
shape={'data': input_shape},
dtype=dtype)
if use_trt:
mod = relay.tensorrt.EnableTrt(mod, params)
assert mod['main'].attrs and mod[
'main'].attrs.Compiler == 'tensorrt'
with relay.build_config(opt_level=2,
disabled_pass={"SimplifyInference"}):
graph, lib, params = relay.build(mod, "cuda")
check_trt_used(graph)
else:
with relay.build_config(opt_level=3):
graph, lib, params = relay.build(mod, "cuda", params=params)
mod = graph_runtime.create(graph, lib, ctx=tvm.gpu(0))
mod.set_input(**params)
# Warmup
for i in range(10):
mod.run(data=i_data)
# Time
times = []
for i in range(num_iteration):
start_time = time.time()
mod.run(data=i_data)
res = mod.get_output(0)
times.append(time.time() - start_time)
latency = 1000.0 * np.mean(times)
print(model, latency)
return latency, res
def get_model(model, ctx, opt):
"""Model initialization."""
kwargs = {'ctx': ctx, 'pretrained': opt.use_pretrained, 'classes': classes}
if model.startswith('resnet'):
kwargs['thumbnail'] = opt.use_thumbnail
elif model.startswith('vgg'):
kwargs['batch_norm'] = opt.batch_norm
net = models.get_model(model, **kwargs)
if opt.resume:
net.load_params(opt.resume)
elif not opt.use_pretrained:
if model in ['alexnet']:
net.initialize(mx.init.Normal())
else:
net.initialize(mx.init.Xavier(magnitude=2))
net.cast(opt.dtype)
return net
def test_models():
all_models = ['resnet18_v1', 'resnet34_v1', 'resnet50_v1', 'resnet101_v1', 'resnet152_v1',
'resnet18_v2', 'resnet34_v2', 'resnet50_v2', 'resnet101_v2', 'resnet152_v2',
'vgg11', 'vgg13', 'vgg16', 'vgg19',
'vgg11_bn', 'vgg13_bn', 'vgg16_bn', 'vgg19_bn',
'alexnet', 'inceptionv3',
'densenet121', 'densenet161', 'densenet169', 'densenet201',
'squeezenet1.0', 'squeezenet1.1']
pretrained_to_test = set(['squeezenet1.1'])
for model_name in all_models:
test_pretrain = model_name in pretrained_to_test
model = get_model(model_name, pretrained=test_pretrain)
data_shape = (2, 3, 224, 224) if 'inception' not in model_name else (2, 3, 299, 299)
eprint('testing forward for %s'%model_name)
print(model)
if not test_pretrain:
model.collect_params().initialize()
model(mx.nd.random_uniform(shape=data_shape)).wait_to_read()
def get_workload(batch_size, num_classes=1000, num_layers=121, dtype="float32"):
"""Get benchmark workload for mobilenet
Parameters
----------
batch_size : int
The batch size used in the model
num_classes : int, optional
Number of classes
num_layers : int, optional
The number of layers
dtype : str, optional
The data type
Returns
-------
net : nnvm.Symbol
The computational graph
params : dict of str to NDArray
The parameters.
"""
import mxnet as mx
from mxnet.gluon.model_zoo.vision import get_model
image_shape = (1, 3, 224, 224)
block = get_model('densenet%d' % num_layers, classes=num_classes, pretrained=False)
data = mx.sym.Variable('data')
sym = block(data)
sym = mx.sym.SoftmaxOutput(sym)
net = _from_mxnet_impl(sym, {})
return create_workload(net, batch_size, image_shape[1:], dtype)
def get_network(name, batch_size):
"""Get the symbol definition and random weight of a network"""
input_shape = (batch_size, 3, 224, 224)
output_shape = (batch_size, 1000)
if "resnet" in name:
n_layer = int(name.split('-')[1])
net, params = nnvm.testing.resnet.get_workload(num_layers=n_layer, batch_size=batch_size)
elif "vgg" in name:
n_layer = int(name.split('-')[1])
net, params = nnvm.testing.vgg.get_workload(num_layers=n_layer, batch_size=batch_size)
elif name == 'mobilenet':
net, params = nnvm.testing.mobilenet.get_workload(batch_size=batch_size)
elif name == 'squeezenet_v1.1':
net, params = nnvm.testing.squeezenet.get_workload(batch_size=batch_size, version='1.1')
elif name == 'inception_v3':
input_shape = (1, 3, 299, 299)
net, params = nnvm.testing.inception_v3.get_workload(batch_size=batch_size)
elif name == 'custom':
# an example for custom network
from nnvm.testing import utils
net = nnvm.sym.Variable('data')
net = nnvm.sym.conv2d(net, channels=4, kernel_size=(3,3), padding=(1,1))
net = nnvm.sym.flatten(net)
net = nnvm.sym.dense(net, units=1000)
net, params = utils.create_workload(net, batch_size, (3, 224, 224))
elif name == 'mxnet':
# an example for mxnet model
from mxnet.gluon.model_zoo.vision import get_model
block = get_model('resnet18_v1', pretrained=True)
net, params = nnvm.frontend.from_mxnet(block)
net = nnvm.sym.softmax(net)
else:
onnx_model = onnx.load_model(
'out/models/resnet50_conv_bs1_0/model.onnx')
net, params = nnvm.frontend.from_onnx(onnx_model)
output_shape = (batch_size, 6, 112, 112)
return net, params, input_shape, output_shape
def get_symbol(network, batch_size, dtype):
image_shape = (3,299,299) if network in ['inception-v3', 'inception-v4'] else (3,224,224)
num_layers = 0
if network == 'inception-resnet-v2':
network = network
elif 'resnet' in network:
num_layers = int(network.split('-')[1])
network = network.split('-')[0]
if 'vgg' in network:
num_layers = int(network.split('-')[1])
network = 'vgg'
if network in ['densenet121', 'squeezenet1.1']:
sym = models.get_model(network)
sym.hybridize()
data = mx.sym.var('data')
sym = sym(data)
sym = mx.sym.SoftmaxOutput(sym, name='softmax')
else:
net = import_module('symbols.'+network)
sym = net.get_symbol(num_classes=1000,
image_shape=','.join([str(i) for i in image_shape]),
num_layers=num_layers,
dtype=dtype)
return (sym, [('data', (batch_size,)+image_shape)])
server = rpc.Server(host=host, port=port, use_popen=True)
######################################################################
# Prepare the Pretrained Model
# ----------------------------
# Back to the host machine, firstly, we need to download a MXNet Gluon
# ResNet model from model zoo, which is pretrained on ImageNet. You
# can found more details about this part at `Compile MXNet Models`
from mxnet.gluon.model_zoo.vision import get_model
from mxnet.gluon.utils import download
from PIL import Image
import numpy as np
# only one line to get the model
block = get_model('resnet18_v1', pretrained=True)
######################################################################
# In order to test our model, here we download an image of cat and
# transform its format.
img_name = 'cat.jpg'
download('https://github.com/dmlc/mxnet.js/blob/master/data/cat.png?raw=true', img_name)
image = Image.open(img_name).resize((224, 224))
def transform_image(image):
image = np.array(image) - np.array([123., 117., 104.])
image /= np.array([58.395, 57.12, 57.375])
image = image.transpose((2, 0, 1))
image = image[np.newaxis, :]
return image
def train(current_host, hosts, num_cpus, num_gpus, channel_input_dirs, model_dir, hyperparameters, **kwargs):
# retrieve the hyperparameters we set in notebook (with some defaults)
batch_size = hyperparameters.get('batch_size', 128)
epochs = hyperparameters.get('epochs', 100)
learning_rate = hyperparameters.get('learning_rate', 0.1)
momentum = hyperparameters.get('momentum', 0.9)
log_interval = hyperparameters.get('log_interval', 1)
wd = hyperparameters.get('wd', 0.0001)
if len(hosts) == 1:
kvstore = 'device' if num_gpus > 0 else 'local'
else:
kvstore = 'dist_device_sync'
ctx = [mx.gpu(i) for i in range(num_gpus)] if num_gpus > 0 else [mx.cpu()]
net = models.get_model('resnet34_v2', ctx=ctx, pretrained=False, classes=10)
batch_size *= max(1, len(ctx))
# load training and validation data
# we use the gluon.data.vision.CIFAR10 class because of its built in pre-processing logic,
# but point it at the location where SageMaker placed the data files, so it doesn't download them again.
part_index = 0
for i, host in enumerate(hosts):
if host == current_host:
part_index = i
break
data_dir = channel_input_dirs['training']
train_data = get_train_data(num_cpus, data_dir, batch_size, (3, 32, 32),
num_parts=len(hosts), part_index=part_index)
test_data = get_test_data(num_cpus, data_dir, batch_size, (3, 32, 32))
# Collect all parameters from net and its children, then initialize them.
net.initialize(mx.init.Xavier(magnitude=2), ctx=ctx)
# Trainer is for updating parameters with gradient.
trainer = gluon.Trainer(net.collect_params(), 'sgd',
optimizer_params={'learning_rate': learning_rate, 'momentum': momentum, 'wd': wd},
kvstore=kvstore)
metric = mx.metric.Accuracy()
loss = gluon.loss.SoftmaxCrossEntropyLoss()
best_accuracy = 0.0
for epoch in range(epochs):
# reset data iterator and metric at begining of epoch.
train_data.reset()
tic = time.time()
metric.reset()
btic = time.time()
for i, batch in enumerate(train_data):
data = gluon.utils.split_and_load(batch.data[0], ctx_list=ctx, batch_axis=0)
label = gluon.utils.split_and_load(batch.label[0], ctx_list=ctx, batch_axis=0)
outputs = []
Ls = []
with ag.record():
for x, y in zip(data, label):
z = net(x)
L = loss(z, y)
# store the loss and do backward after we have done forward
# on all GPUs for better speed on multiple GPUs.
Ls.append(L)
outputs.append(z)
for L in Ls:
L.backward()
trainer.step(batch.data[0].shape[0])
metric.update(label, outputs)
if i % log_interval == 0 and i > 0:
name, acc = metric.get()
logging.info('Epoch [%d] Batch [%d]\tSpeed: %f samples/sec\t%s=%f' %
(epoch, i, batch_size / (time.time() - btic), name, acc))
btic = time.time()
name, acc = metric.get()
logging.info('[Epoch %d] training: %s=%f' % (epoch, name, acc))
logging.info('[Epoch %d] time cost: %f' % (epoch, time.time() - tic))
name, val_acc = test(ctx, net, test_data)
logging.info('[Epoch %d] validation: %s=%f' % (epoch, name, val_acc))
# only save params on primary host
if current_host == hosts[0]:
if val_acc > best_accuracy:
net.save_params('{}/model-{:0>4}.params'.format(model_dir, epoch))
best_accuracy = val_acc
return net
h5f.close()
def save_h5_first(file_name, key, data, l_shape=''):
h5f = h5py.File(file_name, 'w')
if l_shape:
dataset = h5f.create_dataset(key, data=data, maxshape=(None, l_shape))
else:
dataset = h5f.create_dataset(key, data=data, maxshape=(None, ))
h5f.close()
#################### fine-tune gluon
''''''
# resnet152 100352
# inceptionV3 110592
# 获取inceptionV3特征
pretrained_net = models.get_model('inceptionv3', pretrained=True, ctx=ctx)
i = 0
for li_data, label in test_data:
feature = pretrained_net.features(li_data.as_in_context(ctx)) #<NDArray 16x768x12x12 @gpu(0)>
feature = gluon.nn.Flatten()(feature) #<NDArray 16x110592 @gpu(0)>
if(i==0):
save_h5_first('test_inceptionv3.h5', 'features', feature.asnumpy(), feature.shape[1])
#save_h5_first('labels1.h5', 'labels', label.asnumpy())
i+=1
else:
save_h5_append('test_inceptionv3.h5', 'features', feature.asnumpy(), feature.shape[0], feature.shape[1])
#save_h5_append('labels1.h5', 'labels', label.asnumpy(), label.shape[0])
path_imgidx=data_dir+'train.idx',
shuffle=True,
mean=True,
std=True,
rand_crop=1,
min_object_covered=0.95,
max_attempts=200)
batch = trainIter.next()
label, data = batch.label, batch.data
print(label[0].shape,"batchSize, maxObjectNum,classId,xmin,ymin,xmax,ymax,difficult")
print(data[0].shape, "batchSize,C,H,W")
#2--load net symbol
#2.1-load pretrained net (feature part)
pretrained = vision.get_model('resnet18_v1', pretrained=True).features
net = nn.HybridSequential()
for i in range(len(pretrained) - 2):
net.add(pretrained[i])
#not to initialize the pretrained model???
#2.2-add yolo output
class YOLOV1_OUTPUT(nn.HybridBlock):
def __init__(self, numClass, box_per_cell,verbose = False, **kwargs):
super(YOLOV1_OUTPUT, self).__init__(**kwargs)
self.numClass = numClass
self.box_per_cell = box_per_cell
channelNum = box_per_cell * (numClass + 5) #IOU,x,y,w,h
#(x,y) is related to cell
#(w,h) is related to (W,H) of last feature
batch_size, dataset, classes = opt.batch_size, opt.dataset, dataset_classes[opt.dataset]
num_gpus = opt.num_gpus
batch_size *= max(1, num_gpus)
context = [mx.gpu(i) for i in range(num_gpus)] if num_gpus > 0 else [mx.cpu()]
model_name = opt.model
kwargs = {'ctx': context, 'pretrained': opt.use_pretrained, 'classes': classes}
if model_name.startswith('resnet'):
kwargs['thumbnail'] = opt.use_thumbnail
elif model_name.startswith('vgg'):
kwargs['batch_norm'] = opt.batch_norm
net = models.get_model(opt.model, **kwargs)
# get dataset iterators
if dataset == 'mnist':
train_data, val_data = mnist_iterator(batch_size, (1, 32, 32))
elif dataset == 'cifar10':
train_data, val_data = cifar10_iterator(batch_size, (3, 32, 32))
elif dataset == 'imagenet':
if model_name == 'inceptionv3':
train_data, val_data = imagenet_iterator(opt.train_data, opt.val_data,
batch_size, (3, 299, 299))
else:
train_data, val_data = imagenet_iterator(opt.train_data, opt.val_data,
batch_size, (3, 224, 224))
elif dataset == 'dummy':
if model_name == 'inceptionv3':
def get_network(name, batch_size, dtype='float32'):
"""Get the symbol definition and random weight of a network
Parameters
----------
name: str
The name of the network, can be 'resnet-18', 'resnet-50', 'vgg-16', 'inception_v3', 'mobilenet', ...
batch_size: int
batch size
dtype: str
Data type
Returns
-------
net: nnvm.symbol
The NNVM symbol of network definition
params: dict
The random parameters for benchmark
input_shape: tuple
The shape of input tensor
output_shape: tuple
The shape of output tensor
"""
input_shape = (batch_size, 3, 224, 224)
output_shape = (batch_size, 1000)
if name == 'mobilenet':
net, params = nnvm.testing.mobilenet.get_workload(batch_size=batch_size, dtype=dtype)
elif name == 'mobilenet_v2':
net, params = nnvm.testing.mobilenet_v2.get_workload(batch_size=batch_size, dtype=dtype)
elif name == 'inception_v3':
input_shape = (batch_size, 3, 299, 299)
net, params = nnvm.testing.inception_v3.get_workload(batch_size=batch_size, dtype=dtype)
elif "resnet" in name:
n_layer = int(name.split('-')[1])
net, params = nnvm.testing.resnet.get_workload(num_layers=n_layer, batch_size=batch_size, dtype=dtype)
elif "vgg" in name:
n_layer = int(name.split('-')[1])
net, params = nnvm.testing.vgg.get_workload(num_layers=n_layer, batch_size=batch_size, dtype=dtype)
elif "densenet" in name:
n_layer = int(name.split('-')[1])
net, params = nnvm.testing.densenet.get_workload(num_layers=n_layer, batch_size=batch_size, dtype=dtype)
elif "squeezenet" in name:
version = name.split("_v")[1]
net, params = nnvm.testing.squeezenet.get_workload(batch_size=batch_size, version=version, dtype=dtype)
elif name == 'custom':
# an example for custom network
from nnvm.testing import utils
net = nnvm.sym.Variable('data')
net = nnvm.sym.conv2d(net, channels=4, kernel_size=(3,3), padding=(1,1))
net = nnvm.sym.flatten(net)
net = nnvm.sym.dense(net, units=1000)
net, params = utils.create_workload(net, batch_size, (3, 224, 224), dtype=dtype)
elif name == 'mxnet':
# an example for mxnet model
from mxnet.gluon.model_zoo.vision import get_model
block = get_model('resnet18_v1', pretrained=True)
net, params = nnvm.frontend.from_mxnet(block)
net = nnvm.sym.softmax(net)
else:
raise ValueError("Unsupported network: " + name)
return net, params, input_shape, output_shape
def get_nn_model(name):
if "densenet" in name:
return get_model(name, dropout=0)
else:
return get_model(name)
