Support aten::grid_sampler, aten::im2col #8443

delldu · 2021-07-10T13:29:09Z

Dear reviewers,
Help us review this commit, thanks a lot.

wyc-ruiker

I haven't reviewed the specific implementation, but it seems that there are some problems in the unit test.

wyc-ruiker · 2021-07-12T03:13:37Z

tests/python/frontend/pytorch/test_forward.py

+def test_forward_im2col():
+    torch.set_grad_enabled(False)
+
+    class Im2col3x3(Module):


Im2col3x3 is not called in verify_script_model?

Maybe we don't need two separate modules, but use parameters to pass the kernel_size into one Module, so that we can also support more kernel_size in the unit test. @delldu

Im2col3x3 for trace mode, Im2col5x5 for script mode.
ONLY for KISS in unit test. Please feel free comments ^_^.

Modified. Only one class Im2col with constructor, remove script test.

tests/python/frontend/pytorch/test_forward.py

python/tvm/relay/frontend/pytorch.py

wyc-ruiker · 2021-07-12T07:59:20Z

src/relay/op/nn/pad.cc

+  auto kernel_h = tvm::cast(tvm::DataType::Int(32), param->kernel_size[0]);
+  auto kernel_w = tvm::cast(tvm::DataType::Int(32), param->kernel_size[1]);
+  auto dilation_h = tvm::cast(tvm::DataType::Int(32), param->dilation[0]);
+  auto dilation_w = tvm::cast(tvm::DataType::Int(32), param->dilation[1]);
+  auto padding_h = tvm::cast(tvm::DataType::Int(32), param->padding[0]);
+  auto padding_w = tvm::cast(tvm::DataType::Int(32), param->padding[1]);
+  auto stride_h = tvm::cast(tvm::DataType::Int(32), param->stride[0]);
+  auto stride_w = tvm::cast(tvm::DataType::Int(32), param->stride[1]);


It seems that we can use tir::as_const_int?

We checked source code and found that tvm::cast could deal with more complex situation than tir::as_const_int.
Not sure, I am new comer for TVM development, please feel free to figure out my mistake ^-^, thanks.

I don't understand why you need all these casts. kernel_size etc should be already int. Supporting script mode is not an acceptable excuse.

Just copy from PadCompute.

Sorry, let me paste the source code.

bool Im2colRel(const Array<Type>& types, int num_inputs, const Attrs& attrs, const TypeReporter& reporter) { // types: [input, output] ICHECK_EQ(types.size(), 2) << "Expects two types, one for the input and another for the output"; const auto* input = types[0].as<TensorTypeNode>(); if (input == nullptr) return false; if (input->shape.size() != 4) { reporter->GetDiagCtx().EmitFatal(Diagnostic::Error(reporter->GetSpan()) << "Im2lossRel: input data should be 4 dimensions, NxCxHxW."); return false; } const Im2colAttrs* param = attrs.as<Im2colAttrs>(); if (param == nullptr) return false; ICHECK_EQ(param->kernel_size.size(), 2) << "Expects two parameters for kernel height and width"; ICHECK_EQ(param->dilation.size(), 2) << "Expects two parameters for dilation height and width"; ICHECK_EQ(param->padding.size(), 2) << "Expects two parameters for padding height and width"; ICHECK_EQ(param->stride.size(), 2) << "Expects two parameters for stride height and width"; // Calculate output shape auto kernel_h = tvm::cast(tvm::DataType::Int(32), param->kernel_size[0]); auto kernel_w = tvm::cast(tvm::DataType::Int(32), param->kernel_size[1]); auto dilation_h = tvm::cast(tvm::DataType::Int(32), param->dilation[0]); auto dilation_w = tvm::cast(tvm::DataType::Int(32), param->dilation[1]); auto padding_h = tvm::cast(tvm::DataType::Int(32), param->padding[0]); auto padding_w = tvm::cast(tvm::DataType::Int(32), param->padding[1]); auto stride_h = tvm::cast(tvm::DataType::Int(32), param->stride[0]); auto stride_w = tvm::cast(tvm::DataType::Int(32), param->stride[1]); auto dilated_kernel_h = (kernel_h - 1) * dilation_h + 1; auto dilated_kernel_w = (kernel_w - 1) * dilation_w + 1; // Output size after padding auto output_h = (input->shape[2] + 2 * padding_h - dilated_kernel_h) / stride_h + 1; auto output_w = (input->shape[3] + 2 * padding_w - dilated_kernel_w) / stride_w + 1; tvm::Array<tvm::PrimExpr> output_shape; output_shape.push_back(input->shape[0]); // N output_shape.push_back(input->shape[1] * kernel_h * kernel_w); // K output_shape.push_back(output_h * output_w); // L // assign output type reporter->Assign(types[1], TensorType(output_shape, input->dtype)); return true; }

RELAY_REGISTER_OP need Im2col with interface: bool Im2colRel(const Array& types, int num_inputs, const Attrs& attrs, const TypeReporter& reporter);

We must parse attrs for checking they are valid or not;

We must calculate output shape depend on parse results.
Unfortunately, attrs are not integers, so we have to do such stupid things.

src/relay/op/nn/pad.cc

wyc-ruiker · 2021-07-12T08:02:25Z

include/tvm/topi/nn.h

+  auto kernel_h = tvm::cast(tvm::DataType::Int(32), kernel_size[0]);
+  auto kernel_w = tvm::cast(tvm::DataType::Int(32), kernel_size[1]);
+
+  auto dilation_h = tvm::cast(tvm::DataType::Int(32), dilation[0]);
+  auto dilation_w = tvm::cast(tvm::DataType::Int(32), dilation[1]);
+
+  auto padding_h = tvm::cast(tvm::DataType::Int(32), padding[0]);
+  auto padding_w = tvm::cast(tvm::DataType::Int(32), padding[1]);
+
+  auto stride_h = tvm::cast(tvm::DataType::Int(32), stride[0]);
+  auto stride_w = tvm::cast(tvm::DataType::Int(32), stride[1]);


Explained as before. tvm::cast would be a better choice than tir::as_const_int if we run script mode, PyTorch script mode could generate complex expression.

We don't really support script mode. Please don't make your implementation complicated to support script mode.

Just copy from PadCompute.

Please reference on same file about 155~245 lines.

tvm::te::Tensor pad(const tvm::te::Tensor& t, const tvm::Array<tvm::PrimExpr>& pad_before, tvm::Array<tvm::PrimExpr> pad_after = tvm::Array<tvm::PrimExpr>(), PrimExpr pad_value = PrimExpr(), std::string name = "T_pad", std::string tag = kElementWise, std::string pad_mode = "constant", const Array<PrimExpr>* dyn_output_shape = nullptr)

with same function tvm::cast to parsing parameters.
A little more, this is no business with script mode, just for parsing.

@delldu Please drop support for script mode. Just because your model has control flow doesn't mean these ops also need to be scripted. Remember, everyone can use these op conversion. You can do trace and selectively script certain parts of your model using torch.jit._script_if_tracing. This is how MaskRCNN in torchvision is implemented, for example. See https://github.com/pytorch/vision/blob/master/torchvision/models/detection/roi_heads.py#L454

Thanks. It is real good reference. High value !

There is no need to copy old code. kernel_size etc should be constant integers, please remove all casts.

include/tvm/topi/nn.h

tests/python/frontend/pytorch/test_forward.py

masahi · 2021-07-14T05:44:23Z

@delldu I don't see a need to support scripting for these ops, so can you simplify this PR? You don't need to test with verify_script_model.

delldu · 2021-07-15T05:34:16Z

@delldu I don't see a need to support scripting for these ops, so can you simplify this PR? You don't need to test with verify_script_model.

Hello @masahi, @wyc-ruiker, we have simplified test units. Could you please consider accepting this PR ? Or any thing else we need to do ?

wyc-ruiker · 2021-07-15T06:09:25Z

@delldu I don't see a need to support scripting for these ops, so can you simplify this PR? You don't need to test with verify_script_model.

Hello @masahi, @wyc-ruiker, we have simplified test units. Could you please consider accepting this PR ? Or any thing else we need to do ?

Hi @delldu, there's another masahi's comment that hasn't been discussed. #8443 (comment)

masahi · 2021-07-15T06:35:57Z

python/tvm/relay/frontend/pytorch.py

+        Torch F.unfold set kerenl_size, dilation, padding, stride as pairs before calling im2col
+        but it brokern TVM "if condition expression", so please USE torch._C._nn.im2col instead
+        of F.unfold and make sure giving paired parameters. Please reference test_forward_im2col
+        in file tests/python/frontend/pytorch/test_forward.py.


I don't understand this sentences.

From [https://pytorch.org/docs/stable/_modules/torch/nn/functional.html#unfold]:

def unfold( input: Tensor, kernel_size: BroadcastingList2[int], dilation: BroadcastingList2[int] = 1, padding: BroadcastingList2[int] = 0, stride: BroadcastingList2[int] = 1 ) -> Tensor: if input.dim() == 4: msg = "{} must be int or 2-tuple for 4D input" assert_int_or_pair(kernel_size, "kernel_size", msg) assert_int_or_pair(dilation, "dilation", msg) assert_int_or_pair(padding, "padding", msg) assert_int_or_pair(stride, "stride", msg) return torch._C._nn.im2col(input, _pair(kernel_size), _pair(dilation), _pair(padding), _pair(stride)) else: raise NotImplementedError("Input Error: Only 4D input Tensors are supported (got {}D)".format(input.dim()))

You can find:

_pair(kernel_size), _pair(dilation), _pair(padding), _pair(stride)

if input.dim() == 4: ... else raise ...
First thing, we hope user call im2col with pairs parameters. Second, we hope user do not use F.unfold directly for it breaks "if return tensor1 else return tensor2" condition, otherwise, type check will failure.

Removed usefulness comments in source code.

First of all, please fix the typo and grammar error, otherwise the comments do not make sense. I don't understand what is "if condition expression". Please revisit this problem assuming you don't need to support script mode. torch._C._nn.im2col is not supposed to be used by users, we should support the usage of torch.unfold.

masahi · 2021-07-15T06:38:00Z

python/tvm/relay/op/nn/nn.py

+
+    Please reference
+    `https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold`
+


PyTorch specific stuff shouldn't be in this file. Please explain what this function does instead.

If you click "[SOURCE]" link, you will find

def forward(self, input: Tensor) -> Tensor: return F.unfold(input, self.kernel_size, self.dilation, self.padding, self.stride)

torch.nn.Unfold forward will all F.unfold, F.unfold will call im2col. If you call F.unfold or nn.Unfold forward, from_pytorch will give you a exception.

Removed usefulness comments in source code.

Please just remove references to PyTorch and describe what it does. This is not frontend code.

masahi · 2021-07-15T06:43:53Z

tests/python/frontend/pytorch/test_forward.py

+            #
+            # !!! DO NOT USE !!!
+            # F.grid_sample(input, grid, mode='bilinear', padding_mode='zeros', align_corners=True)
+            # for it broken TVM "if conditional expression" in torch script mode


This comment doesn't make sense, remove it.

As F.unfold, F.grid_sample also check parameter with if condition. If you call F.grid_sample, you will get exceptions. This can be reproduced with "F.grid_sample" test. Please kindly reference:
F.grid_sample

Removed usefulness comments in source code.

tests/python/frontend/pytorch/test_forward.py

masahi · 2021-07-15T06:49:40Z

@delldu Please drop support for script mode. Just because your model has control flow doesn't mean these ops also need to be scripted. Remember, everyone can use these op conversion. You can do trace and selectively script certain parts of your model using torch.jit._script_if_tracing. This is how MaskRCNN in torchvision is implemented, for example. See https://github.com/pytorch/vision/blob/master/torchvision/models/detection/roi_heads.py#L454

masahi · 2021-07-15T06:50:38Z

Sorry I closed this PR by mistake.

masahi · 2021-07-15T09:53:41Z

include/tvm/relay/attrs/nn.h

+  Array<IndexExpr> kernel_size;
+  Array<IndexExpr> dilation;
+  Array<IndexExpr> padding;
+  Array<IndexExpr> stride;


IndexExpr should be Integer. If you make this change you can remove all the casts.

Sorry let give a little explain.

Im2Compute must following TVM RELAY_REGISTER_OP interface, so it is

Array<te::Tensor> Im2colCompute(const Attrs& attrs, const Array<te::Tensor>& inputs, const Type& out_type) { const auto* param = attrs.as<Im2colAttrs>(); ICHECK(param != nullptr); return Array<te::Tensor>{ topi::im2col(inputs[0], param->kernel_size, param->dilation, param->padding, param->stride)}; }

We have two choice, 1)parsing parameters in Im2colCompute, 2) parsing in topi::im2col, here we choice 2), that is in topi::im2col, there are no difference. 3) Change Im2colAttrs.

Struct Conv2DAttrs is also using IndexExpr instead of Integer, so Im2colAttrs follow their style, in fact, kernel_size mean kernel_size_h/w, dilation means dilation_h/w etc, not pure Integer.

So RELAY_REGISTER_OP ==> Im2colCompute ==> im2col and Im2colAttrs, we almost have no choice.

I don't see why you can't use Array<Integer>. There is no need to follow Conv2DAttrs, if all attributes in im2col ops are integer constant, there is no point using Array<IndexExpr> and it just makes your implementation unnecessarily complicated.

OK, I see. will modify.

Hello @masahi , @wyc-ruiker ,

First all, thank you help us improve PR. Sounds there is one import issue left, I would give my understanding.

Second, please remember using IndexExpr = ::tvm::PrimExpr;

1. Conclusion
Changing Array<IndexExpr> to Array<Integer> in Im2colAttrs is not good.

2. Function Patch
inline tvm::te::Tensor im2col(const tvm::te::Tensor& data, ...)
{
auto input_b = data->shape[0];
...
auto kernel_h = tvm::cast(tvm::DataType::Int(32), kernel_size[0]);
...
tvm::Arraytvm::PrimExpr output_shape;
output_shape.push_back(N);
...
auto pad_value = tvm::tir::make_const(data->dtype, 0);
auto l = [&](tvm::Arraytvm::tir::Var args) {
...
tvm::PrimExpr s_c = indexdiv(indexdiv(k, kernel_h), kernel_w);
indices.push_back(s_c); // C, source channel
...
return tvm::if_then_else(..., data(indices), pad_value);
};
return tvm::te::compute(output_shape, l, name, tag);
}

3. Cause analysis

Input
The shape of input data is Array<PrimExpr>;
inline tvm::te::Tensor im2col(const tvm::te::Tensor& data

Output
output shape is also Array<PrimExpr>.
TVM_DLL Tensor compute(Array<PrimExpr> shape, ...);

Process
3.1) Many expression in lambda function such as indices, are Array<PrimExpr>, NOT Array<Integer>.
3.2) Some function parameters are PrimExpr.
For an example: prototye of indexmod is:

TVM_DLL PrimExpr indexdiv(PrimExpr a, PrimExpr b, Span span = Span());

If we give Integer parameters to indexdiv, only got compile fatal warning and errors ...

4. Test Result
We try to change Array<IndexExpr> to Array<Integer>, modified related source code,
fighting again and again, using many methods, finally got many and many building(compile) errors.

5. Inference
tvm::cast is good choice than as_const_int for parsing parameters.
The reason comes from their prototype

inline const *int64_t as_const_int(const PrimExpr& x); TVM_DLL PrimExpr cast(const DataType& t, PrimExpr value, Span span = Span());

Following above results, that is why we can't use Array<Integer>.

masahi · 2021-07-15T10:01:03Z

python/tvm/relay/op/nn/_nn.py

@@ -1308,3 +1308,6 @@ def dilate_shape_func(attrs, inputs, _):
 reg.register_shape_func("nn.bias_add", False, elemwise_shape_func)
 reg.register_shape_func("nn.softmax", False, elemwise_shape_func)
 reg.register_shape_func("nn.relu", False, elemwise_shape_func)
+
+
+reg.register_broadcast_schedule("nn.im2col")


Should not be broadcast if you set TOpPattern to kOpaque.

Shall we change kOpaque to kInjective ? Frankly, I do not understand it. please teach me. Thanks.

Let's make it kInjective and use register_injective_schedule here.

Thanks, we have change kOpaque to kInjective, and test passed.

…ol to F.unfold.

delldu · 2021-07-16T14:10:31Z

log.txt

blackkker · 2022-08-16T03:17:00Z

I wonder if the work on aten::im2col is still going on? I got OpNotImplemented error when I run a pytorch model. @masahi @delldu

blackkker · 2022-09-23T07:03:00Z

@masahi @delldu

delldu added 9 commits July 10, 2021 21:24

Add grid_sampler test functions.

4d23916

Support grid_sampler.py

80b2262

Support aten::im2col

fe46820

fix source format with lint.

bc8e101

Fix source format.

23df38b

Fix source format.

1c64c44

fix source format.

6514f74

format include/tvm/relay/attrs/nn.h

bc02b4c

format tests/python/frontend/pytorch/test_forward.py

d086e28

delldu changed the title ~~Support aten::grid_sampler~~ Support aten::grid_sampler, aten::im2col Jul 11, 2021

delldu and others added 3 commits July 12, 2021 01:09

Merge branch 'apache:main' into main

4935e65

Fix bug for aten::Float

303c70a

Test aten::Float for bug fixed

8798175

delldu changed the title ~~Support aten::grid_sampler, aten::im2col~~ Support aten::grid_sampler, aten::im2col, Fix bug aten::Float Jul 11, 2021

wyc-ruiker suggested changes Jul 12, 2021

View reviewed changes

tests/python/frontend/pytorch/test_forward.py Show resolved Hide resolved

Add test_forward_im2col

7a1e24e

wyc-ruiker reviewed Jul 12, 2021

View reviewed changes

delldu added 3 commits July 12, 2021 16:18

Add test_im2col for 3x3 kernel

f25c34b

Add more checks for parameters.

c41c30e

Delete white space in front of line.

c9704d9

wyc-ruiker reviewed Jul 13, 2021

View reviewed changes

tests/python/frontend/pytorch/test_forward.py Outdated Show resolved Hide resolved

Simple test for im2col via removing script mode.

fbf98d2

masahi reviewed Jul 15, 2021

View reviewed changes

tests/python/frontend/pytorch/test_forward.py Outdated Show resolved Hide resolved

masahi reviewed Jul 15, 2021

View reviewed changes

tests/python/frontend/pytorch/test_forward.py Outdated Show resolved Hide resolved

masahi reviewed Jul 15, 2021

View reviewed changes

tests/python/frontend/pytorch/test_forward.py Outdated Show resolved Hide resolved

masahi closed this Jul 15, 2021

masahi reopened this Jul 15, 2021

delldu added 2 commits July 15, 2021 16:13

Remove init constructor from FloatModel.

65bbf30

Dropout script

43c5d01

delldu changed the title ~~Support aten::grid_sampler, aten::im2col, Fix bug aten::Float~~ Support aten::grid_sampler, aten::im2col Jul 15, 2021

masahi reviewed Jul 15, 2021

View reviewed changes

delldu and others added 5 commits July 15, 2021 21:19

Change im2col schudule from broadcast to injective, torch._C_.nn.im2c…

0387aef

…ol to F.unfold.

Reformat source code

7ca82a9

Fix typo mistakes.

5667b94

Remove comments.

97a5df6

Merge branch 'apache:main' into main

8d88fc0

delldu closed this Jul 16, 2021

chengven027 mentioned this pull request Sep 23, 2022

[Torch] A list of missing op conversion in need of help #5133

Closed

12 tasks


		Please reference
		`https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html#torch.nn.Unfold`

Support aten::grid_sampler, aten::im2col #8443

Support aten::grid_sampler, aten::im2col #8443

Conversation

delldu commented Jul 10, 2021

wyc-ruiker left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

delldu Jul 12, 2021 • edited Loading

Choose a reason for hiding this comment

delldu Jul 15, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

delldu Jul 12, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

delldu Jul 15, 2021 • edited Loading

Choose a reason for hiding this comment

wyc-ruiker Jul 12, 2021 • edited Loading

Choose a reason for hiding this comment

delldu Jul 12, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

delldu Jul 15, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

masahi Jul 15, 2021 • edited Loading

Choose a reason for hiding this comment

masahi commented Jul 14, 2021

delldu commented Jul 15, 2021 • edited Loading

wyc-ruiker commented Jul 15, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

delldu Jul 15, 2021 • edited Loading

Choose a reason for hiding this comment

masahi Jul 15, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

masahi commented Jul 15, 2021 • edited Loading

masahi commented Jul 15, 2021

Choose a reason for hiding this comment

delldu Jul 15, 2021 • edited Loading

Choose a reason for hiding this comment

masahi Jul 15, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

delldu Jul 15, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

delldu Jul 15, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

delldu commented Jul 16, 2021

blackkker commented Aug 16, 2022 • edited Loading

blackkker commented Sep 23, 2022

delldu Jul 12, 2021 •

edited

Loading

delldu Jul 15, 2021 •

edited

Loading

delldu Jul 12, 2021 •

edited

Loading

delldu Jul 15, 2021 •

edited

Loading

wyc-ruiker Jul 12, 2021 •

edited

Loading

delldu Jul 12, 2021 •

edited

Loading

delldu Jul 15, 2021 •

edited

Loading

masahi Jul 15, 2021 •

edited

Loading

delldu commented Jul 15, 2021 •

edited

Loading

delldu Jul 15, 2021 •

edited

Loading

masahi Jul 15, 2021 •

edited

Loading

masahi commented Jul 15, 2021 •

edited

Loading

delldu Jul 15, 2021 •

edited

Loading

masahi Jul 15, 2021 •

edited

Loading

delldu Jul 15, 2021 •

edited

Loading

delldu Jul 15, 2021 •

edited

Loading

blackkker commented Aug 16, 2022 •

edited

Loading