import io import os import platform import numpy as np import pytest import megengine.core.tensor.dtype as dtype import megengine.core.tensor.megbrain_graph as G import megengine.functional as F import megengine.module as M import megengine.random as rand from megengine.core._imperative_rt.core2 import apply from megengine.core._wrap import Device from megengine.core.ops import builtin from megengine.device import get_device_count, is_cuda_available from megengine.functional.external import tensorrt_runtime_opr from megengine.jit.tracing import trace from megengine.tensor import Tensor from megengine.utils.comp_graph_tools import GraphInference from megengine.utils.network import Network as Net def check_pygraph_dump(trace_func, inp_data, expect_results, max_err=None): orig_model = io.BytesIO() inp_size = len(inp_data) out_size = len(expect_results) arg_names = ["arg_{}".format(i) for i in range(inp_size)] output_names = ["out_{}".format(i) for i in range(out_size)] trace_func.dump( orig_model, arg_names=arg_names, output_names=output_names, optimize_for_inference=False, ) orig_model.seek(0) net = Net.load(orig_model) file = io.BytesIO() net.dump(file, optimize_for_inference=False) file.seek(0) graph = GraphInference(file) inp_dict = dict([(arg_names[i], inp_data[i].numpy()) for i in range(inp_size)]) results = graph.run(inp_dict=inp_dict) for ind, tensor in enumerate(expect_results): if max_err: np.testing.assert_almost_equal( tensor.numpy(), results[output_names[ind]], max_err ) else: np.testing.assert_equal(tensor.numpy(), results[output_names[ind]]) assert tensor.dtype == results[output_names[ind]].dtype def test_elemwise(): @trace(symbolic=True, capture_as_const=True) def fwd(x, y): z1 = x * y z2 = x + y z3 = z1 / z2 z3 = z3 ** 3 return z3 x = Tensor([1.0, 2.0]) y = Tensor([3.0, 5.0]) result = fwd(x, y) check_pygraph_dump(fwd, [x, y], [result]) def test_reduce(): @trace(symbolic=True, capture_as_const=True) def fwd(data): x = data.sum(axis=2) x = x.mean(axis=1) return x data = Tensor(np.random.random((1, 32, 32))) result = fwd(data) check_pygraph_dump(fwd, [data], [result]) def test_typecvt(): @trace(symbolic=True, capture_as_const=True) def fwd(data): return data.astype(dtype.qint8(0.8)) x = Tensor(np.random.random((2, 3)) * 255) result = fwd(x) check_pygraph_dump(fwd, [x], [result]) def test_matinv(): @trace(symbolic=True, capture_as_const=True) def fwd(data): return F.matinv(data) data = Tensor(np.random.random((5, 5))) result = fwd(data) check_pygraph_dump(fwd, [data], [result]) def test_matmul(): @trace(symbolic=True, capture_as_const=True) def fwd(data1, data2): return F.matmul(data1, data2) data1 = Tensor(np.random.random((32, 64))) data2 = Tensor(np.random.random((64, 16))) result = fwd(data1, data2) check_pygraph_dump(fwd, [data1, data2], [result]) def test_batchmatmul(): @trace(symbolic=True, capture_as_const=True) def fwd(x, y): return F.matmul(x, y) x = Tensor(np.random.random((3, 3, 5))) y = Tensor(np.random.random((3, 5, 3))) result = fwd(x, y) check_pygraph_dump(fwd, [x, y], [result]) def test_dot(): @trace(symbolic=True, capture_as_const=True) def fwd(x, y): return F.dot(x, y) x = Tensor([1.0, 2.0, 3.0]) y = Tensor([3.0, 4.0, 5.0]) result = fwd(x, y) check_pygraph_dump(fwd, [x, y], [result]) def test_svd(): @trace(symbolic=True, capture_as_const=True) def fwd(data): _, out, _ = F.svd(data) return out input = Tensor(np.random.random((1, 1, 3, 3))) result = fwd(input) check_pygraph_dump(fwd, [input], [result]) def test_conv(): conv = M.Conv2d(3, 32, 3) @trace(symbolic=True, capture_as_const=True) def fwd(data): return conv(data) data = Tensor(np.random.random((1, 3, 32, 32))) result = fwd(data) check_pygraph_dump(fwd, [data], [result]) def test_deformable_conv(): if not is_cuda_available(): return conv = M.DeformableConv2d(3, 32, 3) @trace(symbolic=True, capture_as_const=True) def fwd(data, offset, mask): return conv(data, offset, mask) data = Tensor(np.random.random((1, 3, 32, 32))) offset = Tensor(np.ones((32, 3 * 3 * 2, 30, 30)).astype("int32") * 5) mask = Tensor(np.ones((32, 3 * 3, 30, 30)).astype("int32")) out = fwd(data, offset, mask) check_pygraph_dump(fwd, [data, offset, mask], [out]) def test_convtranspose(): deconv = M.ConvTranspose2d(32, 32, 3) @trace(symbolic=True, capture_as_const=True) def fwd(data): return deconv(data) data = Tensor(np.random.random((1, 32, 32, 32))) result = fwd(data) # cu111 has 1e-7 diff check_pygraph_dump(fwd, [data], [result], 5) @pytest.mark.skip(reason="pytest aborted") def test_grouplocal(): n = M.LocalConv2d(3, 32, 32, 32, 3) @trace(symbolic=True, capture_as_const=True) def fwd(data): return n(data) input = Tensor(np.random.random((1, 3, 32, 32))) result = fwd(input) check_pygraph_dump(fwd, [input], [result]) def test_pooling(): @trace(symbolic=True, capture_as_const=True) def fwd(data): out = F.max_pool2d(data, 2, 2) out = F.avg_pool2d(out, 2, 2) return out data = Tensor(np.random.random((1, 3, 64, 64))) result = fwd(data) check_pygraph_dump(fwd, [data], [result]) def test_adaptivepooling(): pool1 = M.AdaptiveMaxPool2d((2, 2)) pool2 = M.AdaptiveAvgPool2d((2, 2)) @trace(symbolic=True, capture_as_const=True) def fwd(data): out = pool1(data) out = pool2(out) return out input = Tensor(np.random.random((1, 3, 32, 32))) result = fwd(input) check_pygraph_dump(fwd, [input], [result]) def test_roipooling(): inp = Tensor(np.random.random((1, 1, 128, 128))) rois = Tensor(np.random.random((4, 5))) @trace(symbolic=True, capture_as_const=True) def fwd(inp, rois): return F.vision.roi_pooling(inp, rois, (2, 2), scale=2.0) output = fwd(inp, rois) check_pygraph_dump(fwd, [inp, rois], [output]) def test_deformable_ps_roi_pooling(): inp = Tensor(np.random.random((1, 256, 64, 64)).astype("float32")) rois = Tensor(np.random.random((1, 5)).astype("float32")) trans = Tensor(np.random.random((24, 2, 7, 7)).astype("float32")) pooled_h = 7 pooled_w = 7 sample_per_part = 4 no_trans = False part_size = 7 spatial_scale = 1.0 / 64 trans_std = 0.1 @trace(symbolic=True, capture_as_const=True) def fwd(inp, rois, trans): y = F.deformable_psroi_pooling( inp, rois, trans, no_trans, part_size, pooled_h, pooled_w, sample_per_part, spatial_scale, trans_std, ) return y result = fwd(inp, rois, trans) check_pygraph_dump(fwd, [inp, rois, trans], [result]) @pytest.mark.skipif( get_device_count("gpu") > 0, reason="does not support int8 when gpu compute capability less than 6.1", ) def test_convbias(): @trace(symbolic=True, capture_as_const=True) def fwd(inp, weight, bias): return F.quantized.conv_bias_activation( inp, weight, bias, dtype=dtype.qint8(scale=1.0), nonlinear_mode="relu" ) inp = Tensor(np.random.random((1, 3, 64, 64)), dtype=dtype.qint8(scale=1.0)) weight = Tensor(np.random.random((32, 3, 3, 3)), dtype=dtype.qint8(scale=1.0)) bias = Tensor(np.random.random((1, 32, 1, 1)), dtype=dtype.qint32(scale=1.0)) result = fwd(inp, weight, bias) check_pygraph_dump(fwd, [inp, weight, bias], [result]) def test_batch_convbias(): if is_cuda_available(): return @trace(symbolic=True, capture_as_const=True) def fwd(inp, weight, bias): return F.quantized.batch_conv_bias_activation( inp, weight, bias, dtype=dtype.qint8(scale=1.0), nonlinear_mode="relu" ) inp = Tensor(np.random.random((1, 3, 64, 64)), dtype=dtype.qint8(scale=1.0)) weight = Tensor(np.random.random((1, 32, 3, 3, 3)), dtype=dtype.qint8(scale=1.0)) bias = Tensor(np.random.random((1, 32, 1, 1)), dtype=dtype.qint32(scale=1.0)) result = fwd(inp, weight, bias) check_pygraph_dump(fwd, [inp, weight, bias], [result]) def test_batchnorm(): bn = M.BatchNorm2d(32) bn.eval() @trace(symbolic=True, capture_as_const=True) def fwd(data): return bn(data) data = Tensor(np.random.random((1, 32, 32, 32))) result = fwd(data) check_pygraph_dump(fwd, [data], [result]) def test_roialign(): inp = Tensor(np.random.randn(1, 1, 128, 128)) rois = Tensor(np.random.random((4, 5))) @trace(symbolic=True, capture_as_const=True) def fwd(inp, rois): return F.vision.roi_align(inp, rois, (2, 2)) output = fwd(inp, rois) check_pygraph_dump(fwd, [inp, rois], [output]) def test_warpperspective(): inp_shape = (1, 1, 4, 4) x = Tensor(np.arange(16, dtype=np.float32).reshape(inp_shape)) M_shape = (1, 3, 3) # M defines a translation: dst(1, 1, h, w) = rst(1, 1, h+1, w+1) M = Tensor( np.array( [[1.0, 0.0, 1.0], [0.0, 1.0, 1.0], [0.0, 0.0, 1.0]], dtype=np.float32 ).reshape(M_shape) ) @trace(symbolic=True, capture_as_const=True) def fwd(x, M): return F.vision.warp_perspective(x, M, (2, 2)) result = fwd(x, M) check_pygraph_dump(fwd, [x, M], [result]) def test_warpaffine(): inp_shape = (1, 3, 3, 3) x = Tensor(np.arange(27, dtype=np.float32).reshape(inp_shape)) weightv = Tensor([[[1.26666667, 0.6, -83.33333333], [-0.33333333, 1, 66.66666667]]]) @trace(symbolic=True, capture_as_const=True) def fwd(x, weightv): return F.vision.warp_affine(x, weightv, (2, 2), border_mode="wrap") outp = fwd(x, weightv) check_pygraph_dump(fwd, [x, weightv], [outp]) def test_remap(): inp_shape = (1, 1, 4, 4) inp = Tensor(np.arange(16, dtype=np.float32).reshape(inp_shape)) map_xy_shape = (1, 2, 2, 2) map_xy = Tensor( np.array( [[[1.0, 0.0], [0.0, 1.0]], [[0.0, 1.0], [0.0, 1.0]]], dtype=np.float32 ).reshape(map_xy_shape) ) @trace(symbolic=True, capture_as_const=True) def fwd(inp, map_xy): return F.vision.remap(inp, map_xy) out = fwd(inp, map_xy) check_pygraph_dump(fwd, [inp, map_xy], [out]) def test_resize(): x = Tensor(np.random.randn(10, 3, 32, 32)) @trace(symbolic=True, capture_as_const=True) def fwd(x): return F.vision.interpolate(x, size=(16, 16), mode="bilinear") out = fwd(x) check_pygraph_dump(fwd, [x], [out]) def test_index_onehot(): src = Tensor([[1.0, 2.0]]) index = Tensor([0]) @trace(symbolic=True, capture_as_const=True) def fwd(src, index): return F.indexing_one_hot(src, index) out = fwd(src, index) check_pygraph_dump(fwd, [src, index], [out]) def test_set_onehot(): x = Tensor(np.arange(1, 4, dtype=np.int32)) @trace(symbolic=True, capture_as_const=True) def fwd(x): return F.one_hot(x, num_classes=4) out = fwd(x) check_pygraph_dump(fwd, [x], [out]) def test_copy(): x = Tensor([1, 2, 3]) @trace(symbolic=True, capture_as_const=True) def fwd(x): return x.to("cpu0:0") o = fwd(x) check_pygraph_dump(fwd, [x], [o]) def test_argsort(): @trace(symbolic=True, capture_as_const=True) def fwd(data): return F.argsort(data, True) data = Tensor([1.0, 2.0, 3.0, 5.0]) result = fwd(data) check_pygraph_dump(fwd, [data], [result]) def test_argmax_min(): @trace(symbolic=True, capture_as_const=True) def fwd(data): return F.argmax(data), F.argmin(data) data = Tensor(np.random.random((10, 10))) result = fwd(data) check_pygraph_dump(fwd, [data], result) def test_condtake(): mask = Tensor(np.array([[True, False], [False, True]], dtype=np.bool_)) x = Tensor(np.array([[1, np.inf], [np.nan, 4]], dtype=np.float32)) @trace(symbolic=True, capture_as_const=True) def fwd(mask, x): v, index = F.cond_take(mask, x) return v, index v, index = fwd(mask, x) check_pygraph_dump(fwd, [mask, x], [v, index]) def test_topk(): x = Tensor(np.array([2, 4, 6, 8, 7, 5, 3, 1], dtype=np.float32)) @trace(symbolic=True, capture_as_const=True) def fwd(x): top, indices = F.topk(x, 5) return top, indices top, indices = fwd(x) check_pygraph_dump(fwd, [x], [top, indices]) def test_random(): @trace(symbolic=True, capture_as_const=True) def fwd(): x = rand.uniform(size=(2, 2)) y = rand.normal(size=(1, 3, 3, 3)) return x, y x, y = fwd() check_pygraph_dump(fwd, [], [x, y]) def test_tensor_gen(): @trace(symbolic=True, capture_as_const=True) def fwd(): a = F.linspace(3, 10, 3, device=Device("xpux").to_c()) b = F.eye(3, device=Device("xpux").to_c()) return a, b a, b = fwd() check_pygraph_dump(fwd, [], [a, b]) def test_getvarshape(): op = builtin.GetVarShape(axis=1) @trace(symbolic=True, capture_as_const=True) def fwd(data): return apply(op, data)[0] data = Tensor(np.random.random((1, 2, 3, 4))) result = fwd(data) check_pygraph_dump(fwd, [data], [result]) def test_concat(): @trace(symbolic=True, capture_as_const=True) def fwd(data1, data2): return F.concat([data1, data2], axis=1) x = Tensor(np.random.random((2, 3))) y = Tensor(np.random.random((2, 5))) result = fwd(x, y) check_pygraph_dump(fwd, [x, y], [result]) def test_broadcast(): inp = Tensor([[1], [2], [3], [4]]) @trace(symbolic=True, capture_as_const=True) def fwd(inp): return F.broadcast_to(inp, (4, 4)) out = fwd(inp) check_pygraph_dump(fwd, [inp], [out]) def test_identity(): @trace(symbolic=True, capture_as_const=True) def fwd(data): return F.copy(data) data = Tensor([1.0, 2.0]) result = fwd(data) check_pygraph_dump(fwd, [data], [result]) @pytest.mark.skip(reason="advance indexing trace error") def test_nms(): x = np.zeros((100, 4)) np.random.seed(42) x[:, :2] = np.random.rand(100, 2) * 20 x[:, 2:] = np.random.rand(100, 2) * 20 + 100 scores = Tensor(np.random.rand(100)) inp = Tensor(x) @trace(symbolic=True, capture_as_const=True) def fwd(inp, scores): return F.nn.nms(inp, scores, iou_thresh=0.7, max_output=3) result = fwd(inp, scores) check_pygraph_dump(fwd, [inp, scores], [result]) def test_dimshuffle(): inp = Tensor([1, 2, 3, 4]) @trace(symbolic=True, capture_as_const=True) def fwd(inp): return inp.T out = fwd(inp) check_pygraph_dump(fwd, [inp], [out]) def test_reshape(): @trace(symbolic=True, capture_as_const=True) def fwd(data): return data.reshape((1, 8)) data = Tensor(np.random.random((1, 2, 2, 2))) result = fwd(data) check_pygraph_dump(fwd, [data], [result]) def test_add_remove_axis(): @trace(symbolic=True, capture_as_const=True) def fwd(data): x = F.expand_dims(data, [0, 0]) y = F.squeeze(x, 0) return y data = Tensor([1.0, 2.0]) result = fwd(data) check_pygraph_dump(fwd, [data], [result]) @pytest.mark.parametrize("mode", ["get", "set", "inc"]) def test_subtensor(mode): items = [[0, True, True, True, False], [1, False, False, False, True]] data = [Tensor(np.random.random((5, 5))), Tensor(np.random.random(2))] if mode == "get": op = builtin.Subtensor(items) data = data[:1] if mode == "set": op = builtin.SetSubtensor(items) if mode == "inc": op = builtin.IncrSubtensor(items) tensors = [Tensor(0), Tensor(4), Tensor(2), Tensor(3)] @trace(symbolic=True, capture_as_const=True) def fwd(*tensors): return apply(op, *tensors)[0] result = fwd(*data, *tensors) check_pygraph_dump(fwd, data + tensors, [result]) @pytest.mark.parametrize("mode", ["get", "set", "inc"]) def test_advance_indexing(mode): items = [[0, False, False, False, True]] tensors = [Tensor([0, 4, 2])] data = [Tensor(np.random.random((5, 5))), Tensor(np.random.random((3, 5)))] if mode == "get": op = builtin.IndexingMultiAxisVec(items) data = data[:1] if mode == "set": op = builtin.IndexingSetMultiAxisVec(items) if mode == "inc": op = builtin.IndexingIncrMultiAxisVec(items) @trace(symbolic=True, capture_as_const=True) def fwd(*tensors): return apply(op, *tensors)[0] result = fwd(*data, *tensors) check_pygraph_dump(fwd, data + tensors, [result]) @pytest.mark.parametrize("mode", ["get", "set", "inc"]) def test_mesh_indexing(mode): items = [[0, True, True, True, False], [1, False, False, False, True]] tensors = [Tensor(0), Tensor(5), Tensor(2), Tensor([1, 3])] data = [Tensor(np.random.random((5, 5))), Tensor(np.random.random((3, 2)))] if mode == "get": op = builtin.IndexingMultiAxisVec(items) data = data[:1] if mode == "set": op = builtin.IndexingSetMultiAxisVec(items) if mode == "inc": op = builtin.IndexingIncrMultiAxisVec(items) @trace(symbolic=True, capture_as_const=True) def fwd(*tensors): return apply(op, *tensors)[0] result = fwd(*data, *tensors) check_pygraph_dump(fwd, data + tensors, [result]) @pytest.mark.parametrize("mode", ["get", "set", "inc"]) def test_batch_mesh_indexing(mode): items = [[1, False, False, False, True], [2, False, False, False, True]] tensors = [Tensor([[0, 2], [0, 2]]), Tensor([[0, 1, 2], [1, 2, 3]])] data = [Tensor(np.random.random((2, 3, 4))), Tensor(np.random.random((2, 2, 3)))] if mode == "get": op = builtin.BatchedMeshIndexing(items) data = data[:1] if mode == "set": op = builtin.BatchedSetMeshIndexing(items) if mode == "inc": op = builtin.BatchedIncrMeshIndexing(items) @trace(symbolic=True, capture_as_const=True) def fwd(*tensors): return apply(op, *tensors)[0] result = fwd(*data, *tensors) check_pygraph_dump(fwd, data + tensors, [result]) @pytest.mark.skip(reason="tmp skip") def test_assert_equal(): g = G.Graph() inp1 = g.make_h2d(dtype=np.float32, device="xpux") inp2 = g.make_h2d(dtype=np.float32, device="xpux") op = builtin.AssertEqual(maxerr=1e-5) out = G.apply_normal_varnode(op, inp1._node, inp2._node)[0] print(out) g.compile(out) file = io.BytesIO() out_model = G.dump_graph([out]) file.write(out_model[0]) file.seek(0) net = Net.load(file) dump_file = io.BytesIO() net.dump(dump_file) dump_file.seek(0) g = GraphInference(dump_file) g.run(np.array([1.0, 2.0]), np.array([1.0, 2.0])) def test_elemwise_multitype(): op = builtin.ElemwiseMultiType(mode="qadd", dtype=dtype.qint32(2.0)) @trace(symbolic=True, capture_as_const=True) def fwd(x, y): return apply(op, x, y)[0] x = Tensor(np.random.random(10) * 10, dtype=dtype.qint8(2.0)) y = Tensor(np.random.random(10) * 10, dtype=dtype.qint8(2.0)) result = fwd(x, y) check_pygraph_dump(fwd, [x, y], [result]) def test_cvtcolor(): inp = np.random.randn(3, 3, 3, 3).astype(np.float32) x = Tensor(inp) @trace(symbolic=True, capture_as_const=True) def fwd(inp): return F.vision.cvt_color(inp, mode="RGB2GRAY") result = fwd(x) check_pygraph_dump(fwd, [x], [result])