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:
- when the target is
cuda
, there will be core dump error. When the target isllvm
, there will not be problems until “inference”. - If I set the
groups
of thenn.ConvTranspose2d
to be 1, there will be no problem(llvm target). However, when I setgroups=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