fix onnx conversion of maxunpool2d
-import warnings
-
import torch
import torch.nn as nn
import torch.nn.functional as F
@@ -36,9 +34,62 @@ class MaxUnpool2dop(Function):
@staticmethod
def symbolic(g, input, indices, kernel_size, stride, padding, output_size):
- warnings.warn(
- 'The definitions of indices are different between Pytorch and ONNX'
- ', so the outputs between Pytorch and ONNX maybe different')
+ # get shape
+ input_shape = g.op('Shape', input)
+ const_0 = g.op('Constant', value_t=torch.tensor(0))
+ const_1 = g.op('Constant', value_t=torch.tensor(1))
+ batch_size = g.op('Gather', input_shape, const_0, axis_i=0)
+ channel = g.op('Gather', input_shape, const_1, axis_i=0)
+
+ # height = (height - 1) * stride + kernel_size
+ height = g.op(
+ 'Gather',
+ input_shape,
+ g.op('Constant', value_t=torch.tensor(2)),
+ axis_i=0)
+ height = g.op('Sub', height, const_1)
+ height = g.op('Mul', height,
+ g.op('Constant', value_t=torch.tensor(stride[1])))
+ height = g.op('Add', height,
+ g.op('Constant', value_t=torch.tensor(kernel_size[1])))
+
+ # width = (width - 1) * stride + kernel_size
+ width = g.op(
+ 'Gather',
+ input_shape,
+ g.op('Constant', value_t=torch.tensor(3)),
+ axis_i=0)
+ width = g.op('Sub', width, const_1)
+ width = g.op('Mul', width,
+ g.op('Constant', value_t=torch.tensor(stride[0])))
+ width = g.op('Add', width,
+ g.op('Constant', value_t=torch.tensor(kernel_size[0])))
+
+ # step of channel
+ channel_step = g.op('Mul', height, width)
+ # step of batch
+ batch_step = g.op('Mul', channel_step, channel)
+
+ # channel offset
+ range_channel = g.op('Range', const_0, channel, const_1)
+ range_channel = g.op(
+ 'Reshape', range_channel,
+ g.op('Constant', value_t=torch.tensor([1, -1, 1, 1])))
+ range_channel = g.op('Mul', range_channel, channel_step)
+ range_channel = g.op('Cast', range_channel, to_i=7) # 7 is int64
+
+ # batch offset
+ range_batch = g.op('Range', const_0, batch_size, const_1)
+ range_batch = g.op(
+ 'Reshape', range_batch,
+ g.op('Constant', value_t=torch.tensor([-1, 1, 1, 1])))
+ range_batch = g.op('Mul', range_batch, batch_step)
+ range_batch = g.op('Cast', range_batch, to_i=7) # 7 is int64
+
+ # update indices
+ indices = g.op('Add', indices, range_channel)
+ indices = g.op('Add', indices, range_batch)
+
return g.op(
'MaxUnpool',
input,
