from tvm.runtime import container
from tvm.runtime import vm as vm_rt
from tvm.relay import testing
from tvm.relay import vm
from tvm.relay.op.contrib import arm_compute_lib
from tvm.contrib.download import download_testdata
from util import load_test_image, build_module, update_lib, get_cpu_op_count, download_model_zoo, parse_options, get_device_arch, get_device_attributes, get_device_type, get_tvm_target
import sys
argv = sys.argv[1:]
device = parse_options(argv)
model_dir = '/mobilenet_v2_1.0_224/'
model_name = 'mobilenet_v2_1.0_224.tflite'
model_dir = download_model_zoo(model_dir, model_name)
tflite_model_file = os.path.join(model_dir, model_name)
tflite_model_buf = open(tflite_model_file, "rb").read()
# Get TFLite model from buffer
try:
import tflite
tflite_model = tflite.Model.GetRootAsModel(tflite_model_buf, 0)
except AttributeError:
import tflite.Model
tflite_model = tflite.Model.Model.GetRootAsModel(tflite_model_buf, 0)
dtype = "float32"
width = 224
height = 224
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)
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 == "inception-v3f":
model_dir = '/inception_v3_2018_04_27/'
model_name = 'inception_v3.tflite'
input_tensor = "input"
input_sshape = (1, 299, 299, 3)
model_dir = download_model_zoo(model_dir, model_name)
tflite_model_file = os.path.join(model_dir, model_name)
tflite_model_buf = open(tflite_model_file, "rb").read()
# Get TFLite model from buffer
try:
import tflite
tflite_model = tflite.Model.GetRootAsModel(tflite_model_buf, 0)
except AttributeError:
import tflite.Model
tflite_model = tflite.Model.Model.GetRootAsModel(
tflite_model_buf, 0)
mod, params = relay.frontend.from_tflite(
tflite_model,
shape_dict={input_tensor: input_sshape},
dtype_dict={input_tensor: "float32"})
desired_layouts = {'nn.conv2d': ['NCHW', 'default']}
seq = tvm.transform.Sequential([
relay.transform.RemoveUnusedFunctions(),
relay.transform.ConvertLayout(desired_layouts)
])
with tvm.transform.PassContext(opt_level=3):
mod = seq(mod)
elif name == "mobilenet-v1f":
model_dir = "/mobilenet_v1_0.5_128/"
model_name = "mobilenet_v1_0.5_128.tflite"
input_tensor = "input"
input_sshape = (1, 128, 128, 3)
model_dir = download_model_zoo(model_dir, model_name)
tflite_model_file = os.path.join(model_dir, model_name)
tflite_model_buf = open(tflite_model_file, "rb").read()
# Get TFLite model from buffer
try:
import tflite
tflite_model = tflite.Model.GetRootAsModel(tflite_model_buf, 0)
except AttributeError:
import tflite.Model
tflite_model = tflite.Model.Model.GetRootAsModel(
tflite_model_buf, 0)
#mod, params = relay.testing.mobilenet.get_workload(batch_size=batch_size)
mod, params = relay.frontend.from_tflite(
tflite_model,
shape_dict={input_tensor: input_sshape},
dtype_dict={input_tensor: "float32"})
desired_layouts = {'nn.conv2d': ['NCHW', 'default']}
seq = tvm.transform.Sequential([
relay.transform.RemoveUnusedFunctions(),
relay.transform.ConvertLayout(desired_layouts)
])
with tvm.transform.PassContext(opt_level=3):
mod = seq(mod)
elif name == "mobilenet-v2f":
model_dir = "/mobilenet_v2_1.0_224/"
model_name = "mobilenet_v2_1.0_224.tflite"
input_tensor = "input"
input_sshape = (1, 224, 224, 3)
model_dir = download_model_zoo(model_dir, model_name)
tflite_model_file = os.path.join(model_dir, model_name)
tflite_model_buf = open(tflite_model_file, "rb").read()
# Get TFLite model from buffer
try:
import tflite
tflite_model = tflite.Model.GetRootAsModel(tflite_model_buf, 0)
except AttributeError:
import tflite.Model
tflite_model = tflite.Model.Model.GetRootAsModel(
tflite_model_buf, 0)
#mod, params = relay.testing.mobilenet.get_workload(batch_size=batch_size)
mod, params = relay.frontend.from_tflite(
tflite_model,
shape_dict={input_tensor: input_sshape},
dtype_dict={input_tensor: "float32"})
desired_layouts = {'nn.conv2d': ['NCHW', 'default']}
seq = tvm.transform.Sequential([
relay.transform.RemoveUnusedFunctions(),
relay.transform.ConvertLayout(desired_layouts)
])
with tvm.transform.PassContext(opt_level=3):
mod = seq(mod)
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 = (batch_size, 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={"data": 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