def test_preprocess(): batch_size = 2 module = Main() data = mge.tensor(np.random.randint(0, 256, size=(batch_size, 3, 48, 160)), dtype=np.float32) traced_module = trace_module(module, {"data": data}) obj = pickle.dumps(traced_module) traced_module = pickle.loads(obj) module = Net(traced_module) module.eval() quad = mge.tensor(np.random.normal(size=(batch_size, 4, 2)), dtype=np.float32) expect = module(data, quad) traced_module = trace_module(module, data, quad) actual = traced_module(data, quad) for i, j in zip(expect, actual): np.testing.assert_array_equal(i, j) func = trace(traced_module, capture_as_const=True) actual = func(data, quad) for i, j in zip(expect, actual): np.testing.assert_array_equal(i, j) model = io.BytesIO() func.dump(model, arg_names=("data", "quad")) model.seek(0) infer_cg = cgtools.GraphInference(model) actual = list( infer_cg.run(inp_dict={ "data": data.numpy(), "quad": quad.numpy() }).values())[0] np.testing.assert_allclose(expect, actual)
def test_preprocess(): module = Main() data = F.ones((1, 14, 8, 8), dtype=np.uint8) traced_module = trace_module(module, data) obj = pickle.dumps(traced_module) traced_module = pickle.loads(obj) module = Net(traced_module) module.eval() idx = F.zeros((1, ), dtype=np.int32) roi = F.ones((1, 2, 2), dtype=np.float32) y = module(data, idx, roi) traced_module = trace_module(module, data, idx, roi) np.testing.assert_array_equal(traced_module(data, idx, roi), y) func = trace(traced_module, capture_as_const=True) np.testing.assert_array_equal(func(data, idx, roi), y) model = io.BytesIO() func.dump(model, arg_names=("data", "idx", "roi")) model.seek(0) infer_cg = cgtools.GraphInference(model) np.testing.assert_allclose( list( infer_cg.run(inp_dict={ "data": data.numpy(), "idx": idx.numpy(), "roi": roi.numpy() }).values())[0], y, atol=1e-6, )
def test_float_func_conv(): class FConvOpr(M.Module): def __init__(self): super().__init__() self.conv = F.conv2d def forward( self, inp, weight, bias=None, stride=(1, 1), padding=(0, 0), dilation=(1, 1), groups=1, ): x = F.conv2d(inp, weight, bias, stride, padding, dilation, groups) return x net = FConvOpr() data = mge.tensor(np.random.random((1, 16, 32, 32))).astype("float32") weight = mge.tensor(np.random.random((32, 16, 2, 2))).astype("float32") traced_module = trace_module(net, data, weight) tm_result = traced_module(data, weight) _test_convert_result([data, weight], traced_module, tm_result, max_err=1e-4)
def test_dump_and_load(): module = MyModule() x = Tensor(np.ones((1, 8, 14, 14))) expect = module(x) traced_module = trace_module(module, x) np.testing.assert_array_equal(expect, traced_module(x)) obj = pickle.dumps(traced_module) new_tm = pickle.loads(obj) _check_id(new_tm) _check_expr_users(new_tm) traced_module.graph._reset_ids() old_nodes = traced_module.graph.nodes().as_list() new_nodes = new_tm.graph.nodes().as_list() old_exprs = traced_module.graph.exprs().as_list() new_exprs = new_tm.graph.exprs().as_list() assert len(old_nodes) == len(new_nodes) for i, j in zip(old_nodes, new_nodes): assert i._name == j._name assert i._qualname == j._qualname assert i._id == j._id assert len(old_exprs) == len(new_exprs) for i, j in zip(old_exprs, new_exprs): assert i._id == j._id np.testing.assert_array_equal(expect, traced_module(x))
def _check_module(build_func: Callable): net = build_func() net.eval() buffer = io.BytesIO() mge.save(net.state_dict(), buffer) buffer.seek(0) inp = Tensor(np.random.random(size=(5, 3, 32, 32))) traced_net = trace_module(build_func(), inp) traced_net.load_state_dict(mge.load(buffer)) _check_param(net, traced_net) buffer.seek(0) traced_net = trace_module(build_func(), inp).flatten() traced_net.load_state_dict(mge.load(buffer)) _check_param(net, traced_net)
def test_id_and_name(): def _check_id(traced_module): _total_ids = traced_module.graph._total_ids node_ids = [n._id for n in traced_module.graph.nodes().as_list()] assert len(set(node_ids)) == len(node_ids) assert max(node_ids) + 1 == _total_ids[0] expr_ids = [n._id for n in traced_module.graph.exprs().as_list()] assert len(set(expr_ids)) == len(expr_ids) assert max(expr_ids) + 1 == _total_ids[1] def _check_name(flatened_module): node_names = [n._name for n in flatened_module.graph.nodes().as_list()] assert len(set(node_names)) == len(node_names) traced_module, x, expect = _init_module() _check_id(traced_module) flattened_module = traced_module.flatten() _check_id(flattened_module) _check_name(flattened_module) # pickle check obj = pickle.dumps(traced_module) traced_module = pickle.loads(obj) Node._set_next_id(159) Expr._set_next_id(1024) graph = traced_module.graph for expr in graph.get_function_by_type(F.relu).as_list(): relu_out = expr.outputs[0] cur_graph = expr.top_graph with cur_graph.insert_exprs(): neg_out = F.neg(relu_out) cur_graph.replace_node({relu_out: neg_out}) cur_graph.compile() _check_id(traced_module) flattened_module = traced_module.flatten() _check_id(flattened_module) _check_name(flattened_module) # check trace TracedModule obj = pickle.dumps(traced_module) traced_module = pickle.loads(obj) module = NewModule(traced_module) traced_module = trace_module(module, x) _check_id(traced_module) flattened_module = traced_module.flatten() _check_id(flattened_module) _check_name(flattened_module)
def test_tensor_method_loader(): class MyModule3(Module): def forward(self, x): return x + 1 m = MyModule3() x = Tensor(np.ones((20))) traced_module = trace_module(m, x) orig_loader_dict = S.TENSORMETHOD_LOADER S.TENSORMETHOD_LOADER = {} @register_tensor_method_loader("__add__") def add_loader(expr): args = list(expr.args) if not isinstance(args[1], TensorNode): args[1] = Tensor(args[1]) node = Constant(args[1], "const").outputs[0] astype_expr = CallMethod(node, "astype") oup = TensorNode( astype_expr, shape=node.shape, dtype=node.dtype, qparams=node.qparams, ) astype_expr.set_args_kwargs(node, expr.inputs[0].dtype) astype_expr.return_val = (oup, ) add_expr = CallMethod(oup, "__add__") add_expr.set_args_kwargs(oup, oup) oup1 = TensorNode( add_expr, shape=oup.shape, dtype=oup.dtype, qparams=node.qparams, ) add_expr.return_val = oup1 args[1] = oup1 expr.set_args_kwargs(*args) obj = pickle.dumps(traced_module) new_module = pickle.loads(obj) _check_expr_users(new_module) _check_id(new_module) result = new_module(x) gt = m(x) assert (isinstance(new_module.graph._exprs[0], Constant) and len(new_module.graph._exprs) == 4) np.testing.assert_equal(result.numpy(), (x + 2).numpy()) S.TENSORMETHOD_LOADER = orig_loader_dict
def test_backward_fold_scale(conv_cls, bn_cls): module = MyModule(conv_cls, bn_cls) module.eval() inp = mge.Tensor(np.random.random((1, 3, 32, 32))) desired = module(inp) traced_net = tm.trace_module(module, inp) traced_net = traced_net.flatten() optimized_net = tm.optimize(traced_net, "BackwardFoldScale") actual = optimized_net(inp) np.testing.assert_allclose(desired=desired, actual=actual, atol=1e-4) # fuse all mul to conv mul_list = optimized_net.graph.get_method_by_type("__mul__").as_list() assert len(mul_list) == 0
def test_jit_trace(): module = MyModule() module.eval() x = F.ones((1, 8, 14, 14)) expect = module(x) traced_module = trace_module(module, x) func = trace(traced_module, capture_as_const=True) np.testing.assert_array_equal(func(x), expect) model = io.BytesIO() func.dump(model) model.seek(0) infer_cg = cgtools.GraphInference(model) np.testing.assert_allclose(list(infer_cg.run(x.numpy()).values())[0], expect, atol=1e-6)
def _check_qualname(net): inp = Tensor(np.random.random(size=(5, 3, 32, 32))) net.eval() traced_net = trace_module(net, inp) base_qualname = traced_net.graph.qualname for node in traced_net.graph.nodes(): qualname = node.qualname qualname = qualname[len(base_qualname) + 1:] if qualname.endswith("]"): qualname = qualname.rsplit(".", 1)[0] if qualname.startswith("["): qualname = "" traced_attr = get_subattr(traced_net, qualname) orig_attr = get_subattr(net, qualname) assert traced_attr is not None assert orig_attr is not None
def test_training_converge(test_traced_module): net = XORNet() if test_traced_module: inp = Tensor(np.random.random((14, 2))) net = trace_module(net, inp) opt = SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4) gm = ad.GradManager().attach(net.parameters()) @trace(symbolic=False) def train(data, label): with gm: pred = net(data) loss = F.nn.cross_entropy(pred, label) gm.backward(loss) optim.clip_grad_norm(net.parameters(), max_norm=0.2, ord=2.0) return loss def infer(data): return net(data) train_dataset = minibatch_generator() losses = [] for data, label in itertools.islice(train_dataset, 2000): data = Tensor(data, dtype=np.float32) label = Tensor(label, dtype=np.int32) opt.clear_grad() loss = train(data, label) optim.clip_grad_value(net.parameters(), lower=-0.1, upper=0.1) opt.step() losses.append(loss.numpy()) assert (np.mean(losses[-100:]) < 0.1), "Final training Loss must be low enough, get {}".format( np.mean(losses[-100:])) ngrid = 10 x = np.linspace(-1.0, 1.0, ngrid) xx, yy = np.meshgrid(x, x) xx = xx.reshape((ngrid * ngrid, 1)) yy = yy.reshape((ngrid * ngrid, 1)) data = mge.tensor(np.concatenate((xx, yy), axis=1).astype(np.float32)) pred = infer(data) precision = calculate_precision(data.numpy(), pred.numpy()) assert precision == 1.0, "Test precision must be high enough, get {}".format( precision)
def test_fuse_bn(conv_cls, bn_cls): module = MyModule(conv_cls, bn_cls) module.eval() inp = mge.Tensor(np.random.random((1, 3, 32, 32))) desired = module(inp) traced_net = tm.trace_module(module, inp) traced_net = traced_net.flatten() optimized_net = tm.optimize(traced_net, "FuseConvBn") actual = optimized_net(inp) np.testing.assert_allclose(desired=desired, actual=actual, atol=1e-4) # fuse all mul to conv bn_list = optimized_net.graph.get_function_by_type(F.batch_norm).as_list() assert len(bn_list) == 0 bn_list = optimized_net.graph.get_module_by_type(M.BatchNorm2d).as_list() assert len(bn_list) == 0
def test_trace_module_2(): class Model(M.Module): def __init__(self): super().__init__() def forward(self, x): out = x.shape out = apply(builtin.Elemwise(mode="ADD"), out, Tensor(1)) return out traced_model = trace_module(Model(), Tensor(([1,]))) assert isinstance(traced_model.graph._exprs[0], Apply) and isinstance( traced_model.graph._exprs[0].opdef, builtin.GetVarShape ) assert isinstance(traced_model.graph._exprs[1], Constant) assert isinstance(traced_model.graph._exprs[2], Apply) and isinstance( traced_model.graph._exprs[2].opdef, builtin.Elemwise ) assert int(traced_model(Tensor([1, 2]))[0]) == 3
def test_functional_loader(): class MyModule2(Module): def forward(self, x, y): return F.conv2d(x, y) m = MyModule2() x = Tensor(np.random.random((1, 3, 32, 32))) y = Tensor(np.random.random((3, 3, 3, 3))) traced_module = trace_module(m, x, y) orig_loader_dict = S.FUNCTIONAL_LOADER S.FUNCTIONAL_LOADER = {} @register_functional_loader(("megengine.functional.nn", "conv2d")) def conv2df_loader(expr): # expr.func = ("megengine.functional.nn","conv2d") kwargs = expr.kwargs orig_weight = expr.named_args["weight"] astype_expr = CallMethod(orig_weight, "astype") oup = TensorNode( astype_expr, shape=orig_weight.shape, dtype=orig_weight.dtype, qparams=orig_weight.qparams, ) astype_expr.set_args_kwargs(orig_weight, expr.named_args["inp"].dtype) astype_expr.return_val = (oup, ) expr.set_arg("weight", oup) obj = pickle.dumps(traced_module) new_module = pickle.loads(obj) _check_expr_users(new_module) _check_id(new_module) result = new_module(x, y) gt = m(x, y) assert (isinstance(new_module.graph._exprs[0], CallMethod) and len(new_module.graph._exprs) == 2) np.testing.assert_equal(result.numpy(), gt.numpy()) S.FUNCTIONAL_LOADER = orig_loader_dict
def test_extra_block(): class PostProcess(M.Module): def forward(self, x): return x * 2 class Net(M.Module): def __init__(self, traced_module): super().__init__() self.post_process = PostProcess() self.traced_module = traced_module def forward(self, x): x = self.traced_module(x) x = self.post_process(x) return x traced_module, x, expect = _init_block() module = Net(traced_module) np.testing.assert_allclose(2 * expect, module(x), atol=1e-6) traced_module = trace_module(module, x) np.testing.assert_allclose(2 * expect, traced_module(x), atol=1e-6)
def test_opdef_loader(): class MyModule1(Module): def forward(self, x, y): op = Elemwise("ADD") return apply(op, x, y)[0] m = MyModule1() x = Tensor(np.ones((20))) y = Tensor(np.ones((20))) traced_module = trace_module(m, x, y) orig_loader_dict = S.OPDEF_LOADER S.OPDEF_LOADER = {} @register_opdef_loader(Elemwise) def add_opdef_loader(expr): if expr.opdef_state["mode"] == "ADD": expr.opdef_state["mode"] = "MUL" node = expr.inputs[1] astype_expr = CallMethod(node, "astype") oup = TensorNode( astype_expr, shape=node.shape, dtype=expr.inputs[0].dtype, qparams=node.qparams, ) astype_expr.set_args_kwargs(node, expr.inputs[0].dtype) astype_expr.return_val = (oup, ) expr.inputs[1] = oup obj = pickle.dumps(traced_module) new_module = pickle.loads(obj) _check_id(new_module) _check_expr_users(new_module) _check_name(new_module.flatten()) assert (isinstance(new_module.graph._exprs[0], CallMethod) and new_module.graph._exprs[1].opdef.mode == "MUL" and len(new_module.graph._exprs) == 2) result = new_module(x, y) np.testing.assert_equal(result.numpy(), x.numpy()) S.OPDEF_LOADER = orig_loader_dict
def test_module_loader(): class MyModule4(Module): def __init__(self): super().__init__() self.conv = M.Conv2d(3, 3, 3) def forward(self, x): return self.conv(x) m = MyModule4() x = Tensor(np.random.random((1, 3, 32, 32))) traced_module = trace_module(m, x) orig_loader_dict = S.MODULE_LOADER S.MODULE_LOADER = {} @register_module_loader(("megengine.module.conv", "Conv2d")) def conv2dm_loader(expr): module = expr.inputs[0].owner args = list(expr.args) orig_inp = args[1] astype_expr = CallMethod(orig_inp, "astype") oup = TensorNode( astype_expr, shape=orig_inp.shape, dtype=orig_inp.dtype, qparams=orig_inp.qparams, ) astype_expr.set_args_kwargs(orig_inp, module.weight.dtype) astype_expr.return_val = (oup, ) args[1] = oup expr.set_args_kwargs(*args) obj = pickle.dumps(traced_module) new_module = pickle.loads(obj) result = new_module(x) gt = m(x) assert (isinstance(new_module.graph._exprs[1], CallMethod) and len(new_module.graph._exprs) == 3) np.testing.assert_equal(result.numpy(), gt.numpy()) S.MODULE_LOADER = orig_loader_dict
def _construct_tm(self): if self.tm is not None: return self.tm all_opr = [] inputs = [] dtypes = [] shapes = [] datas = [] outputs = [] params = [] for opr in self.ir_graph.all_oprs: op_cls = PARAMEXTRACT.get(type(opr), None) if op_cls is not None: params.append(op_cls(opr).extract()) else: params.append({}) all_opr.append(type(opr)) inputs.append([i.name for i in opr.inp_tensors]) dtypes.append([i.dtype for i in opr.inp_tensors]) shapes.append([i.shape for i in opr.inp_tensors]) datas.append([i.np_data for i in opr.inp_tensors]) assert (len(opr.out_tensors) == 1 ), "MegEngine Cannot supports multiple outputs of one Opr" outputs.append([o.name for o in opr.out_tensors]) module = ONNXModule( all_opr, inputs, dtypes, shapes, datas, outputs, params, self.quantizer, self.map_ir_tensor_2_mge_tensor, self.graph_outputs, ) self.tm = tm.trace_module(module, *self.inp_data) return self.tm
def _check_qat_module(qat_net: QATModule): inp = Tensor(np.random.random(size=(5, 3, 32, 32))) traced_net = trace_module(qat_net, inp) for name, qat_module in qat_net.named_modules(): if not isinstance(qat_module, QATModule): continue traced_qat_module = get_subattr(traced_net, name) weight_qparams, act_qparams = get_qparams(qat_module) traced_weight_qparams, traced_act_qparams = get_qparams( traced_qat_module) if weight_qparams: check_qparams(weight_qparams, traced_weight_qparams) if act_qparams: check_qparams(act_qparams, traced_act_qparams) flatten_traced_net = traced_net.flatten() conv0_node = flatten_traced_net.graph.get_node_by_name( "MyModule_block0_conv0").as_unique() conv0_out_node = flatten_traced_net.graph.get_node_by_name( "MyModule_block0_conv0_out").as_unique() assert isinstance(conv0_node.owner, TracedModule) assert conv0_out_node.expr.inputs[0] is conv0_node
def test_shared_module(): class MyModule(M.Module): def __init__(self): super().__init__() self.a = M.Elemwise("ADD") self.b = self.a def forward(self, x, y): z = self.a(x, y) z = self.b(z, y) return z x = Tensor(1) y = Tensor(2) m = MyModule() tm = trace_module(m, x, y) obj = pickle.dumps(tm) load_tm = pickle.loads(obj) _check_expr_users(load_tm) _check_name(load_tm.flatten()) _check_id(load_tm) assert load_tm.a is load_tm.b
def get_traced_module(net, *x): traced_module = trace_module(net, *x) expect = traced_module(*x) return traced_module, expect
def test_trace_module(): enable_expr_checker() x = Tensor(1) m1 = MyModule1() tm1 = trace_module(m1, x) m2 = MyModule2() tm2 = trace_module(m2, x) inp = Tensor(2) gt = m1(inp) output = tm1(inp) for a, b in zip(output, gt): np.testing.assert_equal(a.numpy(), b.numpy()) gt1 = m2(inp) output1 = tm2(inp) for a, b in zip(output1, gt1): np.testing.assert_equal(a.numpy(), b.numpy()) a, b = Tensor(1), Tensor(2) m3 = MyModule3() gt = m3(a, b) tm3 = trace_module(m3, a, b) out = tm3(a, b) np.testing.assert_equal(out.numpy(), gt.numpy()) assert isinstance(tm3.modules.__dict__["0"], M.Elemwise) assert isinstance(tm3.modules.__dict__["2"], TracedModule) assert isinstance(tm3.modules.__dict__["2"].a, M.Elemwise) assert isinstance(tm3.modules.__dict__["3"], M.Elemwise) m4 = MyModule4() tm4 = trace_module(m4, a, b) np.testing.assert_equal(tm4(a, b).numpy(), 3) np.testing.assert_equal(tm4(a, y=b).numpy(), 3) np.testing.assert_equal(tm4(x=a, y=b).numpy(), 3) tm4 = trace_module(m4, a, y=b) np.testing.assert_equal(tm4(a, b).numpy(), 3) np.testing.assert_equal(tm4(a, y=b).numpy(), 3) np.testing.assert_equal(tm4(x=a, y=b).numpy(), 3) tm4 = trace_module(m4, x=a, y=b) np.testing.assert_equal(tm4(a, b).numpy(), 3) np.testing.assert_equal(tm4(a, y=b).numpy(), 3) np.testing.assert_equal(tm4(x=a, y=b).numpy(), 3) tm5 = trace_module(tm4, a, b) np.testing.assert_equal(tm5(a, b).numpy(), 3) np.testing.assert_equal(tm5(a, y=b).numpy(), 3) np.testing.assert_equal(tm5(x=a, y=b).numpy(), 3) tm5 = trace_module(tm4, a, y=b) np.testing.assert_equal(tm5(a, b).numpy(), 3) np.testing.assert_equal(tm5(a, y=b).numpy(), 3) np.testing.assert_equal(tm5(x=a, y=b).numpy(), 3) tm5 = trace_module(tm4, x=a, y=b) np.testing.assert_equal(tm5(a, b).numpy(), 3) np.testing.assert_equal(tm5(a, y=b).numpy(), 3) np.testing.assert_equal(tm5(x=a, y=b).numpy(), 3) assert len(tm4.graph._exprs) == 1 assert isinstance(tm4.graph._exprs[0], CallFunction) class MyModule5(Module): def __init__(self): super().__init__() self.m1 = tm4 def forward(self, x, y): return self.m1(x, y) tm6 = trace_module(MyModule5(), a, b) assert tm6.m1.argspec is None assert tm6.m1._is_top is False
def _init_cls(cls): module = cls() x = F.ones((1, 3, 3, 3)) y = module(x) traced_module = trace_module(module, x) return traced_module, x, y
from test.utils import ConvOpr, dump_mge_model import megengine as mge import numpy as np from megengine.core.tensor import dtype from megengine.quantization.quantize import quantize_qat from megengine.traced_module import trace_module if __name__ == "__main__": net = ConvOpr("normal") traced_module = trace_module(net, mge.tensor(net.data)) mge.save(traced_module, "float_model.tm") dump_mge_model(net, net.data, "float_model") qat_net = quantize_qat(net) inp_dtype = dtype.qint8(16.0 / 128) data = mge.tensor(np.random.random((1, 3, 224, 224))) * 16 data = data.astype(inp_dtype) inp = mge.tensor(dtype.convert_from_qint8(data.numpy())) inp.qparams.scale = mge.tensor(dtype.get_scale(inp_dtype)) inp.qparams.dtype_meta = dtype._builtin_quant_dtypes["qint8"] qat_module = trace_module(qat_net, inp) mge.save(qat_module, "qat_model.tm")