pytorch的pixel unshuffle转到onnx的SpaceToDepth

1. 问题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.

解决方法：
在torch/onnx/symbolic_opset11.py中的205行位置添加以下代码，添加op支持，完成pixel_unshuffle到space2depth的映射。

@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)

2. 问题2，精度对不上

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

解决方法：

在训练的时候采用如下代码，该代码和onnx的排列顺序是一致的。

在转onnx的时候换成F.pixel_unshuffle转，这样子就没有精度误差。

• 训练时：

def forword(self, x):
	feat = space_to_depth(x, scale_factor=self.downscale_factor)
	#feat = F.pixel_unshuffle(x, self.downscale_factor)

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

• 转换时

def forword(self, x):
	#feat = space_to_depth(x, scale_factor=self.downscale_factor)
	feat = F.pixel_unshuffle(x, self.downscale_factor)

问题解决！！

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