【pytorch，onnx，bug解决】pytorch的op：pixel unshuffle转到onnx的op：SpaceToDepth

news2024/12/26 10:43:21

文章目录

- 问题1，转换失败
- 问题1解决
- 问题2，精度对不上
- 问题2解决
- - 训练时
  - 转换时

问题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)

模型可以转换成功，
在这里插入图片描述

但是发现精度对不上。
在这里插入图片描述

问题2，精度对不上

根据某大神分析，

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

这是由于pytorch的F.pixel_unshuffle和onnx的SpaceToDepth 通道排列顺序是不一致的。如下图所示
在这里插入图片描述

问题2解决

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

训练时

![在这里插入图片描述](https://img-blog.csdnimg.cn/873e1afab14e41eabd29e43aa0d0a634.png


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