【pytorch，onnx，bug解决】pytorch的op：pixel unshuffle转到onnx的op：SpaceToDepth_pytorch 图像 unshuffle

问题1， 转换失败

pytorch训练好的模型想要在onnx上部署，但是发现算子F.pixel_unshuffle 不能直接转到onnx 的 SpaceToDepth，
发生以下报错：Exporting the operator pixel_unshuffle to ONNX opset version 11 is not supported. Please feel free to request support or submit a pull request on PyTorch GitHub.

问题1解决

根据pytorch仓库中某大神建议

https://github.com/pytorch/pytorch/pull/53539/commits/decce892b87d725075d2b5dc2031db8610b6be61

在torch/onnx/symbolic_opset11.py中的205行位置添加以下代码，添加op支持。

@parse_args('v', 'i')
def pixel_unshuffle(g, self, downscale_factor):
    rank = sym_help._get_tensor_rank(self)
    if rank is not None and rank != 4:
        return _unimplemented("pixel_unshuffle", "only support 4d input")
    return g.op("SpaceToDepth", self, blocksize_i=downscale_factor)
1
2
3
4
5
6

模型可以转换成功，

但是发现精度对不上。

问题2，精度对不上

根据某大神分析，

https://github.com/onnx/onnx/issues/3739

这是由于pytorch的F.pixel_unshuffle和onnx的SpaceToDepth 通道排列顺序是不一致的。如下图所示

问题2解决

在训练的时候采用如下代码，该代码和onnx的排列顺序是一致的。
在转onnx的时候换成F.pixel_unshuffle转，这样子就没有精度误差。

训练时

def space_to_depth(tensor, scale_factor):
    num, ch, height, width = tensor.shape
    if height % scale_factor != 0 or width % scale_factor != 0:
        raise ValueError('height and widht of tensor must be divisible by '
                         'scale_factor.')

    new_ch = ch * (scale_factor * scale_factor)
    new_height = height // scale_factor
    new_width = width // scale_factor
    tensor = tensor.reshape(num, ch, new_height, scale_factor, new_width, scale_factor)
    tensor = tensor.permute(0, 3, 5, 1, 2, 4)
    tensor = tensor.reshape(num, new_ch, new_height, new_width)

    return tensor
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

转换时

问题解决！！

注意：当然我使用的是opset11，有人是在opset9中修改，具体能不能用我不太清楚，感兴趣的可以去尝试一波。

https://github.com/pytorch/pytorch/pull/72449