From 15f65179d6f0c8932e3d852ecbaaf3cd11bd5f00 Mon Sep 17 00:00:00 2001 From: YJ Shi Date: Fri, 24 Jun 2022 16:54:53 -0700 Subject: [PATCH 1/8] add new ops --- python/tvm/relay/frontend/pytorch.py | 96 ++++++++++++++++++++++++++++ 1 file changed, 96 insertions(+) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index 3e0bf64e4c1c..4b43caac1f28 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -701,6 +701,30 @@ def ones_like(self, inputs, input_types): return out + def new_ones(self, inputs, input_types): + """ + Returns a Tensor of size size filled with 1. By default, the returned Tensor has the same torch.dtype and torch.device as this tensor. + Parameters + size (int...) - a list, tuple, or torch.Size of integers defining the shape of the output tensor. + dtype (torch.dtype, optional) - the desired type of returned tensor. Default: if None, same torch.dtype as this tensor. + device (torch.device, optional) - the desired device of returned tensor. Default: if None, same torch.device as this tensor. + requires_grad (bool, optional) - If autograd should record operations on the returned tensor. Default: False. + + """ + size = inputs[1] + + import torch + + if not isinstance(size, (_expr.Expr, list, torch.Tensor, np.ndarray)): + msg = "Data type %s could not be parsed in ones op" % (type(size)) + raise AssertionError(msg) + + if inputs[2] is not None: + dtype = _convert_dtype_value(inputs[2]) + else: + dtype = input_types[0] + return self.full_impl(size, 1, dtype) + def zeros(self, inputs, input_types): data = inputs[0] @@ -765,6 +789,59 @@ def full_like(self, inputs, input_types): return out + def new_full(self, inputs, input_types): + """ + Returns a Tensor of size size filled with fill_value. + By default, the returned Tensor has the same dtype and device as this tensor. + + Parameters + ---------- + fill_value (scalar) + The number to fill the output tensor with. + dtype (torch.dtype, optional) + The desired type of returned tensor. Default: if None, same torch.dtype as this tensor. + device (torch.device, optional) + The desired device of returned tensor. Default: if None, same torch.device as this tensor. + requires_grad (bool, optional) + If autograd should record operations on the returned tensor. Default: False. + """ + data = inputs[0] + fill_value = inputs[1] + + import torch + + if not isinstance(fill_value, (_expr.Expr, list, torch.Tensor, np.ndarray)): + msg = "Data type %s could not be parsed in full op" % (type(data)) + raise AssertionError(msg) + + if inputs[2] is not None: # dtype given + dtype = _convert_dtype_value(inputs[2]) + else: + # if dtype is None, use the dtype of the input tensor + dtype = self.infer_type(data) + + return self.full_impl(data, fill_value, dtype) + + def fill_(self, inputs, input_types): + data = inputs[0] + fill_value = inputs[1] + dtype = self.infer_type(data) + return self.full_impl(data, fill_value, dtype) + + def pad(self, inputs, input_types): + data = inputs[0] + pad = inputs[1] + mode = inputs[2] + if mode == "replicate" or mode == "circular": + raise ValueError( + "replicate and circular mode for torch.nn.functional.pad are not supported in TVM" + ) + pad_value = inputs[3] + if pad_value == None: + return _op.nn.pad(data, pad, pad_mode=mode) + else: + return _op.nn.pad(data, pad, pad_value=pad_value, pad_mode=mode) + def linspace(self, inputs, input_types): start = inputs[0] stop = inputs[1] @@ -1397,6 +1474,11 @@ def reshape(self, inputs, input_types): new_shape = tmp_shape return _op.transform.reshape(data, new_shape) + def reshape_as(self, inputs, input_types): + data = inputs[0] + new_shape = self.infer_shape(inputs[1]) + return _op.transform.reshape(data, new_shape) + def pixel_shuffle(self, inputs, input_types): data = inputs[0] upscale_factor = inputs[1] @@ -2336,6 +2418,14 @@ def empty(self, inputs, input_types): shape = inputs[0] return _op.zeros(shape, _convert_dtype_value(inputs[1])) + def empty_like(self, inputs, input_types): + """ + Returns an uninitialized tensor with the same size as input. + torch.empty_like(input) is equivalent to torch.empty(input.size(), dtype=input.dtype, layout=input.layout, device=input.device). + """ + shape = self.infer_shape(inputs[0]) + return _op.zeros(shape, _convert_dtype_value(inputs[1])) + def bincount(self, inputs, input_types): data = inputs[0] weights = inputs[1] @@ -3055,8 +3145,12 @@ def create_convert_map(self): "aten::ones_like": self.ones_like, "aten::zeros": self.zeros, "aten::zeros_like": self.zeros_like, + "aten::new_ones": self.new_ones, "aten::full": self.full, "aten::full_like": self.full_like, + "aten::new_full": self.new_full, + "aten::fill_": self.fill_, + "aten::pad": self.pad, "aten::linspace": self.linspace, "aten::reciprocal": self.reciprocal, "aten::repeat": self.repeat, @@ -3121,6 +3215,7 @@ def create_convert_map(self): "aten::size": self.size, "aten::view": self.view, "aten::reshape": self.reshape, + "aten::reshape_as": self.reshape_as, "aten::clone": self.clone, "aten::log_softmax": self.log_softmax, "aten::sigmoid": self.sigmoid, @@ -3239,6 +3334,7 @@ def create_convert_map(self): "aten::tensor": self.identity, # used for example in tensor(1.0) "aten::numel": self.numel, "aten::empty": self.empty, + "aten::empty_like": self.empty_like, "aten::bincount": self.bincount, "aten::scatter_add": self.scatter_add, "aten::__not__": self.logical_not, From ea15144ae8c8b3b369b4bc7d17d084a8fcacc2a9 Mon Sep 17 00:00:00 2001 From: YJ Shi Date: Fri, 24 Jun 2022 17:03:55 -0700 Subject: [PATCH 2/8] fix pad --- python/tvm/relay/frontend/pytorch.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index 4b43caac1f28..219c609bba6a 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -837,10 +837,7 @@ def pad(self, inputs, input_types): "replicate and circular mode for torch.nn.functional.pad are not supported in TVM" ) pad_value = inputs[3] - if pad_value == None: - return _op.nn.pad(data, pad, pad_mode=mode) - else: - return _op.nn.pad(data, pad, pad_value=pad_value, pad_mode=mode) + return _op.nn.pad(data, pad, pad_value=pad_value, pad_mode=mode) def linspace(self, inputs, input_types): start = inputs[0] From a58976e52a7865759f894cc6108e7107402e1f42 Mon Sep 17 00:00:00 2001 From: YJ Shi Date: Fri, 24 Jun 2022 17:08:24 -0700 Subject: [PATCH 3/8] fix pad --- python/tvm/relay/frontend/pytorch.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index 219c609bba6a..5f085e5ed5ce 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -831,12 +831,18 @@ def fill_(self, inputs, input_types): def pad(self, inputs, input_types): data = inputs[0] pad = inputs[1] - mode = inputs[2] + if len(inputs) > 2: + mode = inputs[2] + else: + mode = "constant" + if len(inputs) == 4: + pad_value = inputs[3] + else: + pad_value = 0 if mode == "replicate" or mode == "circular": raise ValueError( "replicate and circular mode for torch.nn.functional.pad are not supported in TVM" ) - pad_value = inputs[3] return _op.nn.pad(data, pad, pad_value=pad_value, pad_mode=mode) def linspace(self, inputs, input_types): From bb5e2dd691396fd8c7415ad8d0e00b617095ab2c Mon Sep 17 00:00:00 2001 From: YJ Shi Date: Mon, 27 Jun 2022 11:26:29 -0700 Subject: [PATCH 4/8] remove pad --- python/tvm/relay/frontend/pytorch.py | 18 ------------------ 1 file changed, 18 deletions(-) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index 5f085e5ed5ce..a7b0d5b06422 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -828,23 +828,6 @@ def fill_(self, inputs, input_types): dtype = self.infer_type(data) return self.full_impl(data, fill_value, dtype) - def pad(self, inputs, input_types): - data = inputs[0] - pad = inputs[1] - if len(inputs) > 2: - mode = inputs[2] - else: - mode = "constant" - if len(inputs) == 4: - pad_value = inputs[3] - else: - pad_value = 0 - if mode == "replicate" or mode == "circular": - raise ValueError( - "replicate and circular mode for torch.nn.functional.pad are not supported in TVM" - ) - return _op.nn.pad(data, pad, pad_value=pad_value, pad_mode=mode) - def linspace(self, inputs, input_types): start = inputs[0] stop = inputs[1] @@ -3153,7 +3136,6 @@ def create_convert_map(self): "aten::full_like": self.full_like, "aten::new_full": self.new_full, "aten::fill_": self.fill_, - "aten::pad": self.pad, "aten::linspace": self.linspace, "aten::reciprocal": self.reciprocal, "aten::repeat": self.repeat, From 74638f964ee64b4e8eae876ac0be9a3cfae36e3e Mon Sep 17 00:00:00 2001 From: YJ Shi Date: Mon, 27 Jun 2022 12:11:11 -0700 Subject: [PATCH 5/8] fix CI --- python/tvm/relay/frontend/pytorch.py | 36 ++++++++++++++++++++-------- 1 file changed, 26 insertions(+), 10 deletions(-) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index a7b0d5b06422..5f2e2f6d3a08 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -703,12 +703,23 @@ def ones_like(self, inputs, input_types): def new_ones(self, inputs, input_types): """ - Returns a Tensor of size size filled with 1. By default, the returned Tensor has the same torch.dtype and torch.device as this tensor. + Returns a Tensor of size size filled with 1. By default, the + returned Tensor has the same torch.dtype and torch.device + as this tensor. + Parameters - size (int...) - a list, tuple, or torch.Size of integers defining the shape of the output tensor. - dtype (torch.dtype, optional) - the desired type of returned tensor. Default: if None, same torch.dtype as this tensor. - device (torch.device, optional) - the desired device of returned tensor. Default: if None, same torch.device as this tensor. - requires_grad (bool, optional) - If autograd should record operations on the returned tensor. Default: False. + size (int...) + a list, tuple, or torch.Size of integers defining the shape of + the output tensor. + dtype (torch.dtype, optional) + the desired type of returned tensor. + Default: if None, same torch.dtype as this tensor. + device (torch.device, optional) + the desired device of returned tensor. + Default: if None, same torch.device as this tensor. + requires_grad (bool, optional) + If autograd should record operations on the returned tensor. + Default: False. """ size = inputs[1] @@ -792,18 +803,22 @@ def full_like(self, inputs, input_types): def new_full(self, inputs, input_types): """ Returns a Tensor of size size filled with fill_value. - By default, the returned Tensor has the same dtype and device as this tensor. + By default, the returned Tensor has the same dtype + and device as this tensor. Parameters ---------- fill_value (scalar) The number to fill the output tensor with. dtype (torch.dtype, optional) - The desired type of returned tensor. Default: if None, same torch.dtype as this tensor. + The desired type of returned tensor. + Default: if None, same torch.dtype as this tensor. device (torch.device, optional) - The desired device of returned tensor. Default: if None, same torch.device as this tensor. + The desired device of returned tensor. + Default: if None, same torch.device as this tensor. requires_grad (bool, optional) - If autograd should record operations on the returned tensor. Default: False. + If autograd should record operations on the returned + tensor. Default: False. """ data = inputs[0] fill_value = inputs[1] @@ -2407,7 +2422,8 @@ def empty(self, inputs, input_types): def empty_like(self, inputs, input_types): """ Returns an uninitialized tensor with the same size as input. - torch.empty_like(input) is equivalent to torch.empty(input.size(), dtype=input.dtype, layout=input.layout, device=input.device). + torch.empty_like(input) is equivalent to torch.empty(input.size(), + dtype=input.dtype, layout=input.layout, device=input.device). """ shape = self.infer_shape(inputs[0]) return _op.zeros(shape, _convert_dtype_value(inputs[1])) From 64daff1cc24a9397c9c956c88fd59f5a8137d5e1 Mon Sep 17 00:00:00 2001 From: YJ Shi Date: Mon, 27 Jun 2022 14:16:42 -0700 Subject: [PATCH 6/8] remove doc --- python/tvm/relay/frontend/pytorch.py | 44 ---------------------------- 1 file changed, 44 deletions(-) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index 5f2e2f6d3a08..c2c22afb2171 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -702,26 +702,6 @@ def ones_like(self, inputs, input_types): return out def new_ones(self, inputs, input_types): - """ - Returns a Tensor of size size filled with 1. By default, the - returned Tensor has the same torch.dtype and torch.device - as this tensor. - - Parameters - size (int...) - a list, tuple, or torch.Size of integers defining the shape of - the output tensor. - dtype (torch.dtype, optional) - the desired type of returned tensor. - Default: if None, same torch.dtype as this tensor. - device (torch.device, optional) - the desired device of returned tensor. - Default: if None, same torch.device as this tensor. - requires_grad (bool, optional) - If autograd should record operations on the returned tensor. - Default: False. - - """ size = inputs[1] import torch @@ -801,25 +781,6 @@ def full_like(self, inputs, input_types): return out def new_full(self, inputs, input_types): - """ - Returns a Tensor of size size filled with fill_value. - By default, the returned Tensor has the same dtype - and device as this tensor. - - Parameters - ---------- - fill_value (scalar) - The number to fill the output tensor with. - dtype (torch.dtype, optional) - The desired type of returned tensor. - Default: if None, same torch.dtype as this tensor. - device (torch.device, optional) - The desired device of returned tensor. - Default: if None, same torch.device as this tensor. - requires_grad (bool, optional) - If autograd should record operations on the returned - tensor. Default: False. - """ data = inputs[0] fill_value = inputs[1] @@ -2420,11 +2381,6 @@ def empty(self, inputs, input_types): return _op.zeros(shape, _convert_dtype_value(inputs[1])) def empty_like(self, inputs, input_types): - """ - Returns an uninitialized tensor with the same size as input. - torch.empty_like(input) is equivalent to torch.empty(input.size(), - dtype=input.dtype, layout=input.layout, device=input.device). - """ shape = self.infer_shape(inputs[0]) return _op.zeros(shape, _convert_dtype_value(inputs[1])) From a4f04407676a691d3d545832ad8cc59cd2f84d1f Mon Sep 17 00:00:00 2001 From: YJ Shi Date: Tue, 28 Jun 2022 17:37:25 -0700 Subject: [PATCH 7/8] fix fill_ --- python/tvm/relay/frontend/pytorch.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index c2c22afb2171..d7765f0e70d3 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -801,8 +801,7 @@ def new_full(self, inputs, input_types): def fill_(self, inputs, input_types): data = inputs[0] fill_value = inputs[1] - dtype = self.infer_type(data) - return self.full_impl(data, fill_value, dtype) + return self.full_impl(self.infer_shape(data), fill_value, input_types[0]) def linspace(self, inputs, input_types): start = inputs[0] From 6eedd049ab56e832b91680ae2a939a6ce90cbe2d Mon Sep 17 00:00:00 2001 From: YJ Shi Date: Wed, 29 Jun 2022 00:43:45 -0700 Subject: [PATCH 8/8] add tests --- python/tvm/relay/frontend/pytorch.py | 21 +++--- tests/python/frontend/pytorch/test_forward.py | 75 +++++++++++++++++++ 2 files changed, 87 insertions(+), 9 deletions(-) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index d7765f0e70d3..2887184b79b8 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -706,7 +706,7 @@ def new_ones(self, inputs, input_types): import torch - if not isinstance(size, (_expr.Expr, list, torch.Tensor, np.ndarray)): + if not isinstance(size, (_expr.Expr, list, tuple, torch.Size, np.ndarray)): msg = "Data type %s could not be parsed in ones op" % (type(size)) raise AssertionError(msg) @@ -781,20 +781,19 @@ def full_like(self, inputs, input_types): return out def new_full(self, inputs, input_types): - data = inputs[0] - fill_value = inputs[1] - + data = inputs[1] + fill_value = inputs[2] import torch - if not isinstance(fill_value, (_expr.Expr, list, torch.Tensor, np.ndarray)): + if not isinstance(data, (_expr.Expr, list, tuple, torch.Size)): msg = "Data type %s could not be parsed in full op" % (type(data)) raise AssertionError(msg) - if inputs[2] is not None: # dtype given - dtype = _convert_dtype_value(inputs[2]) + if inputs[3] is not None: # dtype given + dtype = _convert_dtype_value(inputs[3]) else: # if dtype is None, use the dtype of the input tensor - dtype = self.infer_type(data) + dtype = self.infer_type(input[0]) return self.full_impl(data, fill_value, dtype) @@ -2381,7 +2380,11 @@ def empty(self, inputs, input_types): def empty_like(self, inputs, input_types): shape = self.infer_shape(inputs[0]) - return _op.zeros(shape, _convert_dtype_value(inputs[1])) + if inputs[1] is not None: + dtype = _convert_dtype_value(inputs[1]) + else: + dtype = input_types[0] + return _op.zeros(shape, dtype) def bincount(self, inputs, input_types): data = inputs[0] diff --git a/tests/python/frontend/pytorch/test_forward.py b/tests/python/frontend/pytorch/test_forward.py index e4cb6354c017..3ad4aec77491 100644 --- a/tests/python/frontend/pytorch/test_forward.py +++ b/tests/python/frontend/pytorch/test_forward.py @@ -199,6 +199,28 @@ def visit(op): torch.cuda.empty_cache() +def verify_model_with_input(test_func, input_data, input_dict={}): + baseline_outputs = test_func(*input_data) + trace = torch.jit.trace(test_func, [input.clone() for input in input_data]) + input_names = ["input{}".format(idx) for idx, inp in enumerate(input_data)] + input_shapes = list(zip(input_names, [inp.shape for inp in input_data])) + mod, params = relay.frontend.from_pytorch(trace, input_shapes, {}) + with tvm.transform.PassContext(opt_level=3): + for target in ["llvm", "cuda"]: + if not tvm.runtime.enabled(target): + continue + dev = tvm.device(target, 0) + lib = relay.build(mod, target=target, params=params) + relay_model = graph_executor.GraphModule(lib["default"](dev)) + for name, value in input_dict.items(): + relay_model.set_input(name, value) + relay_model.run() + + compiled_output = relay_model.get_output(0).numpy() + assert_shapes_match(baseline_outputs, compiled_output) + tvm.testing.assert_allclose(baseline_outputs, compiled_output, rtol=1e-5, atol=1e-5) + + # Single operator tests @tvm.testing.uses_gpu def test_forward_pixel_shuffle(): @@ -1275,6 +1297,16 @@ def forward(self, x): verify_model(Reshape3(), input_data=torch.randn(2, 3, 4)) +@tvm.testing.uses_gpu +def test_forward_reshape_as(): + def test_func(input_tensor, other_tensor): + return input_tensor.reshape_as(other_tensor) + + input_data = [torch.rand([2, 1, 10, 1, 10]), torch.rand([2, 1, 10, 10])] + + verify_model_with_input(test_func, input_data, {"input0": input_data[0]}) + + @tvm.testing.uses_gpu def test_flatten(): def _test_flatten(start_dim, end_dim): @@ -2961,6 +2993,17 @@ def forward(self, *args): verify_model(OnesLike3().float().eval(), input_data=input_data) +@tvm.testing.uses_gpu +def test_forward_new_ones(): + torch.set_grad_enabled(False) + input_shape = [1, 3, 10, 10] + + def test_func(input_tensor): + return input_tensor.new_ones([3, 10, 10]) + + verify_model_with_input(test_func, [torch.rand(input_shape).float()]) + + @tvm.testing.uses_gpu def test_forward_zeros(): torch.set_grad_enabled(False) @@ -3034,6 +3077,24 @@ def forward(self, *args): verify_model(FullLike3().float().eval(), input_data=input_data) +@tvm.testing.uses_gpu +def test_forward_new_full(): + torch.set_grad_enabled(False) + input_shape = [1, 3, 10, 10] + + def test_func(input_tensor): + return input_tensor.new_full([2, 3], 1) + + verify_model_with_input(test_func, [torch.rand(input_shape).float()]) + + +def test_forward_fill_(): + def test_func(x): + return x.fill_(3) + + verify_model_with_input(test_func, [torch.rand([1, 3, 10, 10]).float()]) + + @tvm.testing.uses_gpu def test_forward_linspace(): torch.set_grad_enabled(False) @@ -3752,6 +3813,20 @@ def forward(self, data): verify_script_model(Numel(), [(3, 5, 8)], targets) +def test_empty(): + def test_func(): + return torch.empty([1, 3, 10, 10]) + + verify_model_with_input(test_func, []) + + +def test_empty_like(): + def test_func(data): + return torch.empty_like(data) + + verify_model_with_input(test_func, [torch.rand([1, 3, 10, 10]).float()]) + + def test_forward_pretrained_bert_base_uncased(): ###################################################################### # This is an example how to run BERT models using TVM