DepthWise Transposed Conv error

I thought this is some bug, so I opened an issue in the github, but it seems that it is my wrong usage of tvm. So it is better to ask question here:

The short version of my code is like this:

import onnx
import torch
import torch.nn as nn
import tvm
import tvm.relay as relay



class UpSample16(nn.Module):
    def __init__(self, in_chan):
        super(UpSample16, self).__init__()
        self.conv1 = nn.ConvTranspose2d(in_chan,
                           in_chan,
                           kernel_size=(32, 1),
                           stride=(16, 1),
                           padding=(8, 0),
                           groups=in_chan,
                           bias=False)
        self.conv2 = nn.ConvTranspose2d(in_chan,
                           in_chan,
                           kernel_size=(1, 32),
                           stride=(1, 16),
                           padding=(0, 8),
                           groups=in_chan,
                           bias=False)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        return x


input_shape = (1, 3, 768, 768)
shape_dict = {'0': input_shape}
mdpth = './tmp/play.onnx'
opt_level = 3
target = 'cuda'

print('export')
dummy_in = torch.randn(*input_shape)
model = UpSample16(3)
torch.onnx.export(model, dummy_in, mdpth, opset_version=11)

print('load')
model = onnx.load(mdpth)
print('compile')
mod, params = relay.frontend.from_onnx(model, shape_dict, dtype='float32', opset=11)
with relay.build_config(opt_level=opt_level):
    intrp = relay.build_module.create_executor('graph', mod, tvm.gpu(), target)


print('inference')
import cv2
import numpy as np
impth = '000000000139.jpg'
im = cv2.imread(impth)
im = cv2.resize(im, (input_shape[2:][::-1]))
im = (im - np.array([123, 117, 104])) / np.array([58.4, 57.1, 57.4])
im = im.transpose((2, 0, 1))[np.newaxis, :]

out = intrp.evaluate()(tvm.nd.array(im.astype('float32')), **params).asnumpy()

here seems to be two problems:

  1. when the target is cuda , there will be core dump error. When the target is llvm , there will not be problems until “inference”.
  2. If I set the groups of the nn.ConvTranspose2d to be 1, there will be no problem(llvm target). However, when I set groups=in_chan , which means “depthwise-conv-transpose2d”, there will be problem with compiling.

Please take a look at this, if it is brought by my bad usage of tvm, please correct me. Thanks a lot :slight_smile:

1 Like