def check_sharing():
x = relay.var("x", shape=(1, 10))
y = relay.var("y", shape=(1, 10))
z = relay.add(x, y)
func = relay.Function([x, y], z)
x_in = np.ones((1, 10)).astype("float32")
params = {"x": x_in}
graph, lib, params = relay.build(func, target="llvm", params=params)
mod_shared = graph_executor.create(graph, lib, tvm.cpu(0))
mod_shared.load_params(runtime.save_param_dict(params))
num_mods = 10
mods = [
graph_executor.create(graph, lib, tvm.cpu(0))
for _ in range(num_mods)
]
for mod in mods:
mod.share_params(mod_shared, runtime.save_param_dict(params))
a = np.random.uniform(size=(1, 10)).astype("float32")
for mod in mods:
mod.run(y=a)
out = mod.get_output(0, tvm.nd.empty((1, 10)))
np.testing.assert_equal(out.asnumpy(), x_in + a)
# Explicitly delete the shared module and verify correctness.
del mod_shared
for mod in mods:
mod.run(y=a)
out = mod.get_output(0, tvm.nd.empty((1, 10)))
np.testing.assert_equal(out.asnumpy(), x_in + a)
del mod
def build(target_dir):
""" Compiles resnet18 with TVM"""
# Download the pretrained model in MxNet's format.
block = get_model("resnet18_v1", pretrained=True)
shape_dict = {"data": (1, 3, 224, 224)}
mod, params = relay.frontend.from_mxnet(block, shape_dict)
# Add softmax to do classification in last layer.
func = mod["main"]
func = relay.Function(
func.params, relay.nn.softmax(func.body), None, func.type_params, func.attrs
)
target = "llvm"
with tvm.transform.PassContext(opt_level=3):
graph, lib, params = relay.build(func, target, params=params)
# save the model artifacts
deploy_lib = osp.join(target_dir, "deploy_lib.o")
lib.save(deploy_lib)
cc.create_shared(osp.join(target_dir, "deploy_lib.so"), [osp.join(target_dir, "deploy_lib.o")])
with open(osp.join(target_dir, "deploy_graph.json"), "w") as fo:
fo.write(graph)
with open(osp.join(target_dir, "deploy_param.params"), "wb") as fo:
fo.write(runtime.save_param_dict(params))
def test_fp16_build():
dtype = "float16"
dev = tvm.cuda(0)
if dtype == "float16" and not have_fp16(dev.compute_version):
print("skip because gpu does not support fp16")
return
x = relay.var("x", dtype=dtype, shape=(4, 4))
y = relay.var("y", dtype=dtype, shape=(4, 4))
z = x + y
func = relay.Function([x, y], z)
X = tvm.nd.array(np.random.uniform(-1, 1, (4, 4)).astype(dtype), device=dev)
Y = tvm.nd.array(np.random.uniform(-1, 1, (4, 4)).astype(dtype), device=dev)
params = {
"x": X,
"y": Y,
}
# build
g_json, mmod, params = relay.build(func, "cuda", params=params)
# test
rt = tvm.contrib.graph_executor.create(g_json, mmod, dev)
rt.load_params(runtime.save_param_dict(params))
rt.run()
out = rt.get_output(0)
np.testing.assert_allclose(out.numpy(), X.numpy() + Y.numpy(), atol=1e-5, rtol=1e-5)
def build_test_module(opts):
import numpy as np
x = relay.var("x", shape=(10, 5))
y = relay.var("y", shape=(1, 5))
z = relay.add(x, y)
func = relay.Function([x, y], z)
x_data = np.random.rand(10, 5).astype("float32")
y_data = np.random.rand(1, 5).astype("float32")
params = {"y": y_data}
for runtime_name, file_format_str in RUNTIMES.items():
with tvm.transform.PassContext(opt_level=3,
config={"tir.disable_vectorize": True}):
graph, lib, lowered_params = relay.build(
tvm.IRModule.from_expr(func),
f"llvm --runtime={runtime_name} --system-lib",
params=params,
)
build_dir = os.path.abspath(opts.out_dir)
if not os.path.isdir(build_dir):
os.makedirs(build_dir)
ext = "tar" if runtime_name == "c" else "o"
lib_file_name = os.path.join(
build_dir, file_format_str.format(name="test_model", ext=ext))
if runtime_name == "c":
lib.export_library(lib_file_name)
else:
# NOTE: at present, export_libarary will always create _another_ shared object, and you
# can't stably combine two shared objects together (in this case, init_array is not
# populated correctly when you do that). So for now, must continue to use save() with the
# C++ library.
# TODO(areusch): Obliterate runtime.cc and replace with libtvm_runtime.so.
lib.save(lib_file_name)
with open(
os.path.join(
build_dir,
file_format_str.format(name="test_graph", ext="json")),
"w") as f_graph_json:
f_graph_json.write(graph)
with open(
os.path.join(
build_dir,
file_format_str.format(name="test_params", ext="bin")),
"wb") as f_params:
f_params.write(runtime.save_param_dict(lowered_params))
with open(
os.path.join(
build_dir,
file_format_str.format(name="test_data", ext="bin")),
"wb") as fp:
fp.write(x_data.astype(np.float32).tobytes())
x_output = x_data + y_data
with open(
os.path.join(
build_dir,
file_format_str.format(name="test_output", ext="bin")),
"wb") as fp:
fp.write(x_output.astype(np.float32).tobytes())
def import_graphdef(
name,
batch_size,
seq_len,
save_relay=True,
relay_file="model.json",
relay_params="model.params",
):
abs_path = os.path.dirname(os.path.abspath(__file__))
shape_dict = {"input_1": (batch_size, seq_len)}
relay_file = ("%s_%d_%d_%s" %
(name, batch_size, seq_len, relay_file)).replace("/", "_")
relay_params = ("%s_%d_%d_%s" %
(name, batch_size, seq_len, relay_params)).replace(
"/", "_")
if os.path.exists(os.path.join(abs_path, relay_file)) and os.path.exists(
os.path.join(abs_path, relay_params)):
with open(os.path.join(abs_path, relay_file), "r") as fi:
mod = tvm.ir.load_json(fi.read())
with open(os.path.join(abs_path, relay_params), "rb") as fi:
params = relay.load_param_dict(fi.read())
else:
graph_def = download_model(name, batch_size, seq_len)
mod, params = relay.frontend.from_tensorflow(graph_def,
shape=shape_dict)
if save_relay:
with open(os.path.join(abs_path, relay_file), "w") as fo:
fo.write(tvm.ir.save_json(mod))
with open(os.path.join(abs_path, relay_params), "wb") as fo:
fo.write(runtime.save_param_dict(params))
return mod, dict(params.items()), shape_dict
def build_graph_lib(model_file, opt_level):
"""Compiles the pre-trained model with TVM"""
out_dir = os.path.join(sys.path[0], "../lib")
if not os.path.exists(out_dir):
os.makedirs(out_dir)
# Compile the relay mod
mod, params = _get_mod_and_params(model_file)
target = "llvm -mtriple=wasm32-unknown-unknown -mattr=+simd128 --system-lib"
with tvm.transform.PassContext(opt_level=opt_level):
graph_json, lib, params = relay.build(mod,
target=target,
params=params)
# Save the model artifacts to obj_file
obj_file = os.path.join(out_dir, "graph.o")
lib.save(obj_file)
# Run llvm-ar to archive obj_file into lib_file
lib_file = os.path.join(out_dir, "libgraph_wasm32.a")
cmds = [os.environ.get("LLVM_AR", "llvm-ar-10"), "rcs", lib_file, obj_file]
subprocess.run(cmds)
with open(os.path.join(out_dir, "graph.json"), "w") as f_graph:
f_graph.write(graph_json)
with open(os.path.join(out_dir, "graph.params"), "wb") as f_params:
f_params.write(runtime.save_param_dict(params))
def export_library(self, directory_path):
"""Export the pipeline executor into disk files.
Parameters
----------
directory_path : str
Export the files to this directory.
"""
if not self.pipeline_mods:
raise RuntimeError("The pipeline executor has not been initialized.")
# Check if the directory_path exists.
if not os.path.exists(directory_path):
raise RuntimeError("The directory {directory_path} does not exist.")
# Create an load configuration.
load_config_file_name = "{}/load_config".format(directory_path)
pipeline_config_file_name = "{}/pipeline_config".format(directory_path)
config = {}
config["load_config"] = load_config_file_name
config["pipeline_config"] = pipeline_config_file_name
load_config = []
# Export the library, JSON, and parameter into files, then export these files path
# into a configuration file.
for lib_index in self.pipeline_mods:
mconfig = {}
mconfig["mod_idx"] = lib_index
mconfig["lib_name"] = "{}/lib{}.so".format(directory_path, lib_index)
mconfig["json_name"] = "{}/json{}".format(directory_path, lib_index)
mconfig["params_name"] = "{}/params{}".format(directory_path, lib_index)
mconfig["dev"] = "{},{}".format(
self.pipeline_mods[lib_index]["dev"].device_type,
self.pipeline_mods[lib_index]["dev"].device_id,
)
# Get the graph, lib, and parameters from GraphExecutorFactoryModule.
lib = self.pipeline_mods[lib_index]["lib"]
# Export the lib, graph, and parameters to disk.
lib.export_library(mconfig["lib_name"])
with open(mconfig["json_name"], "w") as file_handle:
file_handle.write(lib.graph_json)
with open(mconfig["params_name"], "wb") as file_handle:
file_handle.write(runtime.save_param_dict(lib.params))
load_config.append(mconfig)
with open(load_config_file_name, "w") as file_handle:
json.dump(load_config, file_handle)
with open(pipeline_config_file_name, "w") as file_handle:
json.dump(self.mods_config, file_handle)
config_file_name = "{}/config".format(directory_path)
with open(config_file_name, "w") as file_handle:
json.dump(config, file_handle)
return config_file_name
def test_save_load():
x = np.ones((10, 2)).astype("float32")
y = np.ones((1, 2, 3)).astype("float32")
params = {"x": x, "y": y}
param_bytes = runtime.save_param_dict(params)
assert isinstance(param_bytes, bytearray)
param2 = relay.load_param_dict(param_bytes)
assert len(param2) == 2
np.testing.assert_equal(param2["x"].numpy(), x)
np.testing.assert_equal(param2["y"].numpy(), y)
def main():
dshape = (32, 16)
net = _get_model(dshape)
mod, params = testing.create_workload(net)
graph, lib, params = relay.build(mod, "llvm", params=params)
with open(osp.join(CWD, "graph.json"), "w") as f_resnet:
f_resnet.write(graph)
with open(osp.join(CWD, "graph.params"), "wb") as f_params:
f_params.write(runtime.save_param_dict(params))
def verify_rpc_gpu_remove_package_params(obj_format):
if not tvm.testing.device_enabled("cuda"):
print("Skip because cuda is not enabled")
return
mod, params = relay.testing.synthetic.get_workload()
with relay.build_config(opt_level=3):
complied_graph_lib = relay.build_module.build(mod,
"cuda",
params=params)
from tvm.contrib import utils
temp = utils.tempdir()
if obj_format == ".so":
file_name = "deploy_lib.so"
else:
assert obj_format == ".tar"
file_name = "deploy_lib.tar"
path_lib = temp.relpath(file_name)
complied_graph_lib_no_params = complied_graph_lib["remove_params"]()
complied_graph_lib_no_params.export_library(path_lib)
path_params = temp.relpath("deploy_param.params")
with open(path_params, "wb") as fo:
fo.write(runtime.save_param_dict(complied_graph_lib.get_params()))
from tvm import rpc
remote = rpc.LocalSession()
remote.upload(path_lib)
loaded_lib = remote.load_module(path_lib)
data = np.random.uniform(-1, 1,
size=input_shape(mod)).astype("float32")
dev = remote.cuda()
# raw api
gmod = loaded_lib["default"](dev)
set_input = gmod["set_input"]
run = gmod["run"]
get_output = gmod["get_output"]
load_params = gmod["load_params"]
loaded_params = bytearray(open(path_params, "rb").read())
set_input("data", tvm.nd.array(data, device=dev))
load_params(loaded_params)
run()
out = get_output(0).numpy()
tvm.testing.assert_allclose(out, verify(data), atol=1e-5)
# graph executor wrapper
gmod = graph_executor.GraphModule(loaded_lib["default"](dev))
loaded_params = bytearray(open(path_params, "rb").read())
gmod.set_input("data", data)
gmod.load_params(loaded_params)
gmod.run()
out = gmod.get_output(0).numpy()
tvm.testing.assert_allclose(out, verify(data), atol=1e-5)
def build_graph_lib(opt_level):
"""Compiles the pre-trained model with TVM"""
out_dir = os.path.join(sys.path[0], "../lib")
if not os.path.exists(out_dir):
os.makedirs(out_dir)
# Follow the tutorial to download and compile the model
model_path = download_testdata(model_url,
"resnet50-v2-7.onnx",
module="onnx")
onnx_model = onnx.load(model_path)
img_url = "https://s3.amazonaws.com/model-server/inputs/kitten.jpg"
img_path = download_testdata(img_url, "imagenet_cat.png", module="data")
# Resize it to 224x224
resized_image = Image.open(img_path).resize((224, 224))
img_data = np.asarray(resized_image).astype("float32")
# Our input image is in HWC layout while ONNX expects CHW input, so convert the array
img_data = np.transpose(img_data, (2, 0, 1))
# Normalize according to the ImageNet input specification
imagenet_mean = np.array([0.485, 0.456, 0.406]).reshape((3, 1, 1))
imagenet_stddev = np.array([0.229, 0.224, 0.225]).reshape((3, 1, 1))
norm_img_data = (img_data / 255 - imagenet_mean) / imagenet_stddev
# Add the batch dimension, as we are expecting 4-dimensional input: NCHW.
img_data = np.expand_dims(norm_img_data, axis=0)
input_name = "data"
shape_dict = {input_name: img_data.shape}
mod, params = relay.frontend.from_onnx(onnx_model, shape_dict)
target = "llvm -mtriple=wasm32-unknown-unknown -mattr=+simd128 --system-lib"
with tvm.transform.PassContext(opt_level=opt_level):
factory = relay.build(mod, target=target, params=params)
# Save the model artifacts to obj_file
obj_file = os.path.join(out_dir, "graph.o")
factory.get_lib().save(obj_file)
# Run llvm-ar to archive obj_file into lib_file
lib_file = os.path.join(out_dir, "libgraph_wasm32.a")
cmds = [os.environ.get("LLVM_AR", "llvm-ar-10"), "rcs", lib_file, obj_file]
subprocess.run(cmds)
# Save the json and params
with open(os.path.join(out_dir, "graph.json"), "w") as f_graph:
f_graph.write(factory.get_graph_json())
with open(os.path.join(out_dir, "graph.params"), "wb") as f_params:
f_params.write(runtime.save_param_dict(factory.get_params()))
def main():
dshape = (4, 8)
net = _get_model(dshape)
mod, params = testing.create_workload(net)
graph, lib, params = relay.build(mod, "llvm --system-lib", params=params)
out_dir = sys.argv[1]
lib.save(osp.join(sys.argv[1], "graph.o"))
with open(osp.join(out_dir, "graph.json"), "w") as f_resnet:
f_resnet.write(graph)
with open(osp.join(out_dir, "graph.params"), "wb") as f_params:
f_params.write(runtime.save_param_dict(params))
def test_ndarray_reflection():
# Make two `NDArrayWrapper`s that point to the same underlying array.
np_array = np.random.uniform(size=(10, 2)).astype("float32")
tvm_array = tvm.nd.array(np_array)
param_dict = {"x": tvm_array, "y": tvm_array}
assert param_dict["x"].same_as(param_dict["y"])
# Serialize then deserialize `param_dict`.
deser_param_dict = relay.load_param_dict(
runtime.save_param_dict(param_dict))
# Make sure the data matches the original data and `x` and `y` contain the same data.
np.testing.assert_equal(deser_param_dict["x"].numpy(), tvm_array.numpy())
# Make sure `x` and `y` contain the same data.
np.testing.assert_equal(deser_param_dict["x"].numpy(),
deser_param_dict["y"].numpy())
def save_module(module_path, graph, lib, params, cross=None):
"""
Create a tarball containing the generated TVM graph,
exported library and parameters
Parameters
----------
module_path : str
path to the target tar.gz file to be created,
including the file name
graph : str
A JSON-serialized TVM execution graph.
lib : tvm.module.Module
A TVM module containing the compiled functions.
params : dict
The parameters (weights) for the TVM module.
cross : str or callable object, optional
Function that performs the actual compilation
"""
lib_name = "mod.so"
graph_name = "mod.json"
param_name = "mod.params"
temp = utils.tempdir()
path_lib = temp.relpath(lib_name)
if not cross:
logger.debug("exporting library to %s", path_lib)
lib.export_library(path_lib)
else:
logger.debug("exporting library to %s , using cross compiler %s",
path_lib, cross)
lib.export_library(path_lib, cc.cross_compiler(cross))
with open(temp.relpath(graph_name), "w") as graph_file:
logger.debug("writing graph to file to %s", graph_file.name)
graph_file.write(graph)
with open(temp.relpath(param_name), "wb") as params_file:
logger.debug("writing params to file to %s", params_file.name)
params_file.write(runtime.save_param_dict(params))
logger.debug("saving module as tar file to %s", module_path)
with tarfile.open(module_path, "w") as tar:
tar.add(path_lib, lib_name)
tar.add(temp.relpath(graph_name), graph_name)
tar.add(temp.relpath(param_name), param_name)
def test_load_unexpected_params():
# Test whether graph_executor.load_params works if parameters
# are provided that are not an expected input.
mod = tvm.IRModule()
params = {}
x = relay.var("x", shape=(1, 10))
y = relay.var("y", shape=(1, 10))
z = relay.add(x, y)
mod["main"] = relay.Function([x, y], z)
graph_module = relay.build(mod, target="llvm", params=params)
rt_mod = tvm.contrib.graph_executor.create(graph_module.get_json(),
graph_module.get_lib(),
tvm.cpu(0))
new_params = graph_module.get_params()
new_params.update({"y_unknown": np.ones((1, )).astype("float32")})
rt_mod.load_params(runtime.save_param_dict(new_params))
def build_module(opts):
dshape = (1, 3, 224, 224)
from mxnet.gluon.model_zoo.vision import get_model
block = get_model("mobilenet0.25", pretrained=True)
shape_dict = {"data": dshape}
mod, params = relay.frontend.from_mxnet(block, shape_dict)
func = mod["main"]
func = relay.Function(func.params, relay.nn.softmax(func.body), None,
func.type_params, func.attrs)
for runtime_name, file_format_str in RUNTIMES.items():
with tvm.transform.PassContext(opt_level=3,
config={"tir.disable_vectorize": True}):
graph, lib, params = relay.build(
func,
f"llvm --runtime={runtime_name} --system-lib",
params=params)
build_dir = os.path.abspath(opts.out_dir)
if not os.path.isdir(build_dir):
os.makedirs(build_dir)
ext = "tar" if runtime_name == "c" else "o"
lib_file_name = os.path.join(
build_dir, file_format_str.format(name="model", ext=ext))
if runtime_name == "c":
lib.export_library(lib_file_name)
else:
# NOTE: at present, export_libarary will always create _another_ shared object, and you
# can't stably combine two shared objects together (in this case, init_array is not
# populated correctly when you do that). So for now, must continue to use save() with the
# C++ library.
# TODO(areusch): Obliterate runtime.cc and replace with libtvm_runtime.so.
lib.save(lib_file_name)
with open(
os.path.join(build_dir,
file_format_str.format(name="graph", ext="json")),
"w") as f_graph_json:
f_graph_json.write(graph)
with open(
os.path.join(build_dir,
file_format_str.format(name="params", ext="bin")),
"wb") as f_params:
f_params.write(runtime.save_param_dict(params))
def verify_graph_executor(remote, target, shape, dtype):
x = relay.var("x")
y = relay.const(1)
z = relay.add(x, y)
func = relay.Function([x], z)
x_in = np.ones(shape).astype(dtype)
params = {"x": x_in}
graph, lib, params = relay.build(func, target=target, params=params)
temp = utils.tempdir()
path_dso = temp.relpath("dev_lib.o")
lib.save(path_dso)
remote.upload(path_dso)
lib = remote.load_module("dev_lib.o")
dev = remote.cpu(0)
mod = graph_executor.create(graph, lib, dev)
mod.load_params(runtime.save_param_dict(params))
mod.run()
out = mod.get_output(0, tvm.nd.empty(shape, dtype=dtype, device=dev))
tvm.testing.assert_allclose(x_in + 1, out.numpy())
def test_load_params_with_constants_in_ext_codegen():
# After binding params and partitioning graph_module.get_params()
# might contain parameters that are not an graph runtime input but
# for example constants in external function.
y_in = np.ones((1,)).astype("float32")
params = {"y": y_in}
mod = tvm.IRModule()
x = relay.var("x", shape=(1, 10))
y = relay.var("y", shape=(1,))
xcb = compiler_begin(x, "ccompiler")
ycb = compiler_begin(y, "ccompiler")
z = relay.add(xcb, ycb)
zce = compiler_end(z, "ccompiler")
mod["main"] = relay.Function([x, y], zce)
mod["main"] = bind_params_by_name(mod["main"], params)
mod = transform.PartitionGraph()(mod)
graph_module = relay.build(mod, target="llvm", params=params)
lib = update_lib(graph_module.get_lib())
rt_mod = tvm.contrib.graph_runtime.create(graph_module.get_json(), lib, tvm.cpu(0))
rt_mod.load_params(runtime.save_param_dict(graph_module.get_params()))
input_shape = (1, 224, 224, 3)
input_dtype = "float32"
# parse TFLite model and convert into Relay computation graph
mod, params = relay.frontend.from_tflite(
tflite_model, shape_dict={input_tensor: input_shape}, dtype_dict={input_tensor: input_dtype}
)
#############
# Compilation
# -----------
target = "llvm"
# Build with Relay
with transform.PassContext(opt_level=3):
graph, lib, params = relay.build_module.build(mod, target, params=params)
###############################################
# Save the graph, lib and parameters into files
# ---------------------------------------------
lib.export_library("./mobilenet.so")
print("lib export succeefully")
with open("./mobilenet.json", "w") as fo:
fo.write(graph)
with open("./mobilenet.params", "wb") as fo:
fo.write(runtime.save_param_dict(params))
def _serialize_params(
params: Optional[Dict[str, NDArray]]) -> Optional[bytearray]:
if params is None:
return None
return save_param_dict(params)
def _get_network(
args: Tuple[str, List[int]]
) -> Tuple[IRModule, bytearray, Tuple[str, List[int], str]]:
name: str
input_shape: List[int]
name, input_shape = args
mod: IRModule
if name in [
"resnet_18",
"resnet_50",
"wide_resnet_50",
"resnext_50",
"mobilenet_v2",
"mobilenet_v3",
"inception_v3",
"densenet_121",
"resnet3d_18",
"vgg_16",
]:
import torch # type: ignore
from torchvision import models # type: ignore
if name in ["resnet_18", "resnet_50"]:
model = getattr(models, name.replace("_", ""))(pretrained=False)
elif name == "wide_resnet_50":
model = getattr(models, "wide_resnet50_2")(pretrained=False)
elif name == "resnext_50":
model = getattr(models, "resnext50_32x4d")(pretrained=False)
elif name == "mobilenet_v2":
model = getattr(models, name)(pretrained=False)
elif name == "mobilenet_v3":
model = getattr(models, name + "_large")(pretrained=False)
elif name == "inception_v3":
model = getattr(models, name)(pretrained=False, aux_logits=False)
elif name == "densenet_121":
model = getattr(models, name.replace("_", ""))(pretrained=False)
elif name == "resnet3d_18":
model = models.video.r3d_18(pretrained=False)
elif name == "vgg_16":
model = getattr(models, name.replace("_", ""))(pretrained=False)
dtype = "float32"
input_data = torch.randn(input_shape).type( # pylint: disable=no-member
{
"float32": torch.float32, # pylint: disable=no-member
}[dtype])
scripted_model = torch.jit.trace(model, input_data).eval()
input_name = "input0"
shape_list = [(input_name, input_shape)]
mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)
with tvm.transform.PassContext(opt_level=3):
mod = tvm.transform.Sequential([
relay.transform.RemoveUnusedFunctions(),
relay.transform.ConvertLayout({
"nn.conv2d": ["NHWC", "default"],
"nn.conv3d": ["NDHWC", "default"],
"nn.max_pool2d": ["NHWC", "default"],
"nn.avg_pool2d": ["NHWC", "default"],
}),
])(mod)
inputs = (input_name, input_shape, dtype)
elif name in ["bert_tiny", "bert_base", "bert_medium", "bert_large"]:
os.environ["TOKENIZERS_PARALLELISM"] = "false"
# pip3 install transformers==3.5 torch==1.7
import torch # type: ignore
import transformers # type: ignore
config_dict = {
"bert_tiny":
transformers.BertConfig(
num_hidden_layers=6,
hidden_size=512,
intermediate_size=2048,
num_attention_heads=8,
return_dict=False,
),
"bert_base":
transformers.BertConfig(
num_hidden_layers=12,
hidden_size=768,
intermediate_size=3072,
num_attention_heads=12,
return_dict=False,
),
"bert_medium":
transformers.BertConfig(
num_hidden_layers=12,
hidden_size=1024,
intermediate_size=4096,
num_attention_heads=16,
return_dict=False,
),
"bert_large":
transformers.BertConfig(
num_hidden_layers=24,
hidden_size=1024,
intermediate_size=4096,
num_attention_heads=16,
return_dict=False,
),
}
configuration = config_dict[name]
model = transformers.BertModel(configuration)
input_name = "input_ids"
input_dtype = "int64"
a = torch.randint(10000, input_shape) # pylint: disable=no-member
model.eval()
scripted_model = torch.jit.trace(model, [a], strict=False)
input_name = "input_ids"
shape_list = [(input_name, input_shape)]
mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)
mod = relay.transform.FastMath()(mod)
mod = relay.transform.CombineParallelBatchMatmul()(mod)
inputs = (input_name, input_shape, input_dtype)
elif name == "dcgan":
output_shape = input_shape
batch_size = output_shape[0]
oshape = output_shape[1:]
mod, params = relay.testing.dcgan.get_workload(
batch_size=batch_size,
oshape=oshape,
layout="NHWC",
)
inputs = ("data", [100], "float32")
else:
raise ValueError("Invalid name: " + name)
params_bytearray: bytearray = save_param_dict(params)
return mod, params_bytearray, inputs
