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