def test_dump_backward_graph(): x0 = tensor(np.random.randn(3, 4)) x1 = tensor(np.random.randn(3, 4)) gm = GradManager().attach(x0) @trace(symbolic=True, capture_as_const=True) def f(x0, x1): with gm: y = x0 * x1 gm.backward(y, F.ones_like(y)) dx0 = x0.grad return y, dx0 y, dx0 = f(x0, x1) np.testing.assert_equal(dx0.numpy(), x1) file = io.BytesIO() f.dump(file, optimize_for_inference=False) file.seek(0) infer_cg = cgtools.GraphInference(file) results = list((infer_cg.run(x0, x1)).values()) np.testing.assert_equal(results[0], y) np.testing.assert_equal(results[1], dx0)
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_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_graph_function(): class Net(M.Module): def forward(self, a, b): return a - b, a * b net = Net() @trace(symbolic=True, capture_as_const=True) def function(a, b, *, net=None): return net(a, b) a = np.array([1, 2, 3]) b = np.array([3]) x, y = function(megengine.tensor(a), megengine.tensor(b), net=net) file = io.BytesIO() function.dump( file, arg_names=["a", "b"], output_names=["x", "y"], optimize_for_inference=False, ) file.seek(0) graph = cgtools.GraphInference(file) results = graph.run(inp_dict={"a": a, "b": b}) np.testing.assert_equal(x.numpy(), results["x"]) np.testing.assert_equal(y.numpy(), results["y"]) results = graph.run(a, inp_dict={"b": b}) np.testing.assert_equal(x.numpy(), results["x"]) np.testing.assert_equal(y.numpy(), results["y"]) results = graph.run(a, b) np.testing.assert_equal(x.numpy(), results["x"]) np.testing.assert_equal(y.numpy(), results["y"]) file.seek(0) graph1 = cgtools.GraphInference(file, outputs=["x"]) results = graph1.run(inp_dict={"a": a, "b": b}) np.testing.assert_equal(x.numpy(), results["x"]) assert "y" not in results
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 test_subtensor_when_shape_invalid(): @jit.trace(symbolic=True, capture_as_const=True) def fun(inp): shape = inp.shape H = shape[-1] NH = H * 8 + 4 arr = F.arange(4, NH, 8) arr_shape = arr.shape return arr_shape[0] inp = rand.uniform(size=[1, 3, 224, 224]) fun(inp) with NamedTemporaryFile() as f: fun.dump(f.name, arg_names=["data"], optimize_for_inference=True) inp = rand.uniform(size=[1, 3, 512, 512]) net = cgtools.GraphInference(f.name) net.run(inp_dict={"data": inp})
def test_capture_dump(): a = tensor([2]) @trace(symbolic=True, capture_as_const=True) def f(x): return x * a x = tensor([3]) y = f(x).numpy() for i in range(3): np.testing.assert_equal(f(x).numpy(), y) file = io.BytesIO() f.dump(file) file.seek(0) infer_cg = cgtools.GraphInference(file) result = list((infer_cg.run(x)).values())[0] np.testing.assert_equal(result[0], y)
def test_dump(): @trace(symbolic=True, capture_as_const=True) def f(a, b): return a + b a = tensor([2]) b = tensor([4]) y = f(a, b).numpy() for i in range(3): np.testing.assert_equal(f(a, b).numpy(), y) file = io.BytesIO() dump_info = f.dump(file) assert dump_info.nr_opr == 3 np.testing.assert_equal(dump_info.inputs, ["arg_0", "arg_1"]) np.testing.assert_equal(dump_info.outputs, ["ADD(arg_0,arg_1)[4]"]) file.seek(0) infer_cg = cgtools.GraphInference(file) result = list((infer_cg.run(a, b)).values())[0] np.testing.assert_equal(result[0], y)
def run_index(index): inp, outp = get_param(cases, index) inp_tensor = [make_tensor(inpi, network) for inpi in inp] if test_trace and not network: copied_inp = inp_tensor.copy() for symbolic in [False, True]: traced_func = trace(symbolic=symbolic)(func) for _ in range(3): traced_results = traced_func(*copied_inp, **kwargs) check_results(traced_results, outp) dumped_func = trace(symbolic=True, capture_as_const=True)(func) dumped_results = dumped_func(*copied_inp, **kwargs) check_results(dumped_results, outp) file = io.BytesIO() dump_info = dumped_func.dump(file) file.seek(0) # arg_name has pattern arg_xxx, xxx is int value def take_number(arg_name): return int(arg_name.split("_")[-1]) input_names = dump_info[4] inps_np = [i.numpy() for i in copied_inp] input_names.sort(key=take_number) inp_dict = dict(zip(input_names, inps_np)) infer_cg = cgtools.GraphInference(file) # assume #outputs == 1 loaded_results = list(infer_cg.run(inp_dict=inp_dict).values())[0] check_results(loaded_results, outp, check_shape=False) # scalar info lost results = func(*inp_tensor, **kwargs) check_results(results, outp, check_shape=(network is None))
def test_dump(dump_format): @trace(symbolic=True, capture_as_const=True) def f(a, b): return a + b # prevent from remaining scope from exception test AutoNaming.clear() a = tensor([2]) b = tensor([4]) y = f(a, b).numpy() for i in range(3): np.testing.assert_equal(f(a, b).numpy(), y) file = io.BytesIO() dump_info = f.dump(file, dump_format=dump_format) assert dump_info.nr_opr == 3 np.testing.assert_equal(dump_info.inputs, ["arg_0", "arg_1"]) np.testing.assert_equal(dump_info.outputs, ["ADD"]) file.seek(0) infer_cg = cgtools.GraphInference(file) result = list((infer_cg.run(a, b)).values())[0] np.testing.assert_equal(result[0], y)
def infer_mge(x, file): infer_cg = cgtools.GraphInference(file + ".mge") # pylint: disable=no-member y = list(infer_cg.run(x).values())[0] print(y.mean()) return y
def test_quantize_batchmatmul_activation(): batch = 4 in_features = 8 out_features = 4 class TestNet(Module): def __init__(self, bias): super().__init__() self.quant = QuantStub() self.dequant = DequantStub() self.batch_mm = BatchMatMulActivation(batch, in_features, out_features, bias=bias) def forward(self, inp): out = self.quant(inp) out = self.batch_mm(out) out = expand_dims(out, -1) out = self.dequant(out) return out inputs = tensor( np.random.randn(batch, in_features, out_features).astype(np.float32)) for bias in (True, False): net = TestNet(bias) net.train() qat_net = quantize_qat(net, inplace=False) disable_fake_quant(qat_net) normal_outputs = net(inputs) qat_outputs = qat_net(inputs) np.testing.assert_allclose(normal_outputs.numpy(), qat_outputs.numpy()) net.eval() normal_outputs = net(inputs) qat_net.eval() qat_outputs = qat_net(inputs) np.testing.assert_allclose(normal_outputs.numpy(), qat_outputs.numpy()) enable_fake_quant(qat_net) qat_outputs = qat_net(inputs) qnet = quantize(qat_net, inplace=False) qnet.eval() quantize_outputs = qnet(inputs) np.testing.assert_allclose(qat_outputs.numpy(), quantize_outputs.numpy(), atol=1e-6) @jit.trace(capture_as_const=True) def f(x): qnet.eval() return qnet(x) f(inputs) file = io.BytesIO() f.dump(file, enable_nchw4=True) file.seek(0) infer_cg = cgtools.GraphInference(file)[0] dumped_outputs = list(infer_cg.run(inputs.numpy()).values())[0] np.testing.assert_allclose(quantize_outputs.numpy(), dumped_outputs, atol=1e-6)
traced_func = trace(symbolic=symbolic)(func) for _ in range(3): traced_results = traced_func(*copied_inp, **kwargs) check_results(traced_results, outp) dumped_func = trace(symbolic=True, capture_as_const=True)(func) dumped_results = dumped_func(*copied_inp, **kwargs) check_results(dumped_results, outp) file = io.BytesIO() dump_info = dumped_func.dump(file) file.seek(0) # arg_name has pattern arg_xxx, xxx is int value def take_number(arg_name): return int(arg_name.split("_")[-1]) input_names = dump_info[4] inps_np = [i.numpy() for i in copied_inp] input_names.sort(key=take_number) inp_dict = dict(zip(input_names, inps_np)) infer_cg = cgtools.GraphInference(file) # assume #outputs == 1 loaded_results = list(infer_cg.run(inp_dict=inp_dict).values())[0] check_results(loaded_results, outp) results = func(*inp_tensor, **kwargs) check_results(results, outp)