如何使用 PyTorch 正确实现数据重组？

Posted 2023-03-27

【中文标题】如何使用 PyTorch 正确实现数据重组？

【英文标题】：How to properly implement data reorganization using PyTorch?

【发布时间】：2020-03-29 07:36:13

【问题描述】：

这将是一个很长的帖子，提前抱歉......

我正在研究一种去噪算法，我的目标是：

使用 PyTorch 设计/训练模型 将 PyTorch 模型转换为 CoreML 模型

去噪算法由以下3部分组成：

“下采样”+ 噪声级别图

一个常规的卷积网络

“上采样”

第一部分的想法很简单，但不太容易解释。例如，给定一个输入彩色图像和一个表示图像噪声标准偏差的输入值“sigma”。 “下采样”部分实际上是空间到深度。简而言之，对于给定的通道和 2x2 像素的子集，空间到深度创建由 4 个通道组成的单个像素。通道数乘以 4，而高度和宽度除以 2。数据被简单地重新组织。 噪声水平图包括创建 3 个包含标准偏差值的通道，以便卷积网络知道如何正确地对输入图像进行去噪。 使用一些代码可能会更清楚：

def downsample_and_noise_map(input, sigma):

    # Input tensor size (batch, channels, height, width)
    in_n, in_c, in_h, in_w = input.size()

    # Output tensor size
    out_h = in_h // 2
    out_w = in_w // 2
    sigma_c = in_c      # nb of channels of the standard deviation tensor
    image_c = in_c * 4  # nb of channels of the image tensor

    # Standard deviation tensor
    output_sigma = sigma.view(1, 1, 1, 1).repeat(in_n, sigma_c, out_h, out_w)

    # Image tensor
    output_image = torch.zeros((in_n, image_c, out_h, out_w))
    output_image[:, 0::4, :, :] = input[:, :, 0::2, 0::2]
    output_image[:, 1::4, :, :] = input[:, :, 0::2, 1::2]
    output_image[:, 2::4, :, :] = input[:, :, 1::2, 0::2]
    output_image[:, 3::4, :, :] = input[:, :, 1::2, 1::2]

    # Concatenate standard deviation and image tensors
    return torch.cat((output_sigma, output_image), dim=1)

这个函数随后在模型的forward 函数中作为第一步调用：

def forward(self, x, sigma):
    x = downsample_and_noise_map(x, sigma)
    x = self.convnet(x)
    x = upsample(x)
    return x

让我们考虑一个大小为 1x3x100x100（PyTorch 标准：批处理、通道、高度、宽度）的输入张量和 0.1 的 sigma 值。输出张量具有以下属性：

张量的形状是 1x15x50x50 通道 0、1 和 2 的张量值都等于 sigma = 0.1 通道 3、4、5、6 的张量值由通道 0 的输入图像值组成 通道 7、8、9、10 的张量值由通道 1 的输入图像值组成 通道 11、12、13、14 的张量值由通道 2 的输入图像值组成

如果这段代码不够清晰，我可以发一个更幼稚的版本。

上采样部分是下采样部分的倒数。

我能够使用这个函数在 PyTorch 中进行训练和测试。

然后，我尝试使用 ONNX 作为中间步骤将模型转换为 CoreML。 转换为 ONNX 会生成“TracerWarning”。从 ONNX 到 CoreML 的转换失败（TypeError：1.0 的类型为 numpy.float64，但预期为以下之一：int、long）。问题来自下采样 + 噪声级别图（也来自上采样）。

当我移除下采样 + 噪声级别图和上采样层时，我能够非常轻松地转换为 ONNX 和 CoreML，因为只剩下一个简单的卷积网络。这意味着我有一个解决问题的方法：在移动端使用 2 个着色器实现这 2 个图层。但我对这个解决方案不满意，因为我希望我的模型包含所有层 ^^

在考虑在这里写帖子之前，我爬网寻找答案，我能够使用reshape 和permute 编写之前函数的更好版本。这个版本去掉了所有的ONNX警告，但是CoreML转换还是失败了……

def downsample_and_noise_map(input, sigma):

    # Input image size
    in_n, in_c, in_h, in_w = input.size()

    # Output tensor size
    out_n = in_n
    out_h = in_h // 2
    out_w = in_w // 2

    # Create standard deviation tensor
    output_sigma = sigma.view(out_n, 1, 1, 1).repeat(out_n, in_c, out_h, out_w)

    # Split RGB channels
    channels_rgb = torch.split(input, 1, dim=1)

    # Reshape (space-to-depth) each image channel
    channels_reshaped = []
    for channel in channels_rgb:
        channel = channel.reshape(1, out_h, 2, out_w, 2)
        channel = channel.permute(2, 4, 0, 1, 3)
        channel = channel.reshape(1, 4, out_h, out_w)
        channels_reshaped.append(channel)

    # Concatenate all reshaped image channels together
    output_image = torch.cat(channels_reshaped, dim=1)

    # Concatenate standard deviation and image tensors
    output = torch.cat([output_sigma, output_image], dim=1)

    return output

所以这是我的（一些）问题：

在模型中实现 downsample_and_noise_map 等函数的首选 PyTorch 方法是什么？ 同样的问题，但何时转换为 ONNX，然后转换为 CoreML 是等式的一部分？ PyTorch -> ONNX -> CoreML 仍然是为 ios 生产部署模型的最佳途径吗？

感谢您的帮助（和您的耐心）^^

【问题讨论】：

你也可以发布上采样代码吗？这对我的测试很有用。

我完成了我的答案（抱歉多次编辑）。请让我知道它是怎么回事，如果它回答了您的问题，请接受作为答案:)

展开 for 循环有帮助吗？不幸的是，没有直接的 PyTorch --> Core ML 转换。必须通过 ONNX 意味着转换可能会在更多地方出错（例如，在 PyTorch --> ONNX 实际上被破坏之前，我遇到过问题）。

@jodag 感谢您的回答。请参阅下面的答案:)

@MatthijsHollemans Jodag 提出的解决方案似乎在使用 for 循环时也有效。命令sigma.view(1, 1, 1, 1).repeat(in_n, sigma_c, out_h, out_w) 是一个似乎没有被转换很好接受的例子......您关于两步转换风险的评论是非常正确的。我希望我有一个 PyTorch 到 CoreML 的转换器，甚至更好：一个适用于 iOS 和 android 的 GPU 加速 PyTorch 库:)

【参考方案1】：

免责声明我不熟悉 CoreML 或部署到 iOS，但我确实有通过 ONNX 在 TensorRT 和 OpenVINO 中部署 PyTorch 模型的经验。

我在部署到其他框架时遇到的主要问题是，像切片和重复张量这样的操作在其他框架中的支持往往有限。通常我们可以构造等效的 conv 或 transpose-conv 操作来实现所需的行为。

为了确保我们不会导出用于构造 conv 权重的逻辑，我将权重初始化与权重应用分开。这使得 ONNX 导出更加直接，因为它所看到的只是应用了一些常量张量。

class DownsampleAndNoiseMap():
    def __init__(self):
        self.initialized = False
        self.weight = None
        self.zeros = None

    def init_weights(self, input):
        with torch.no_grad():
            in_n, in_c, in_h, in_w = input.size()

            out_h = int(in_h // 2)
            out_w = int(in_w // 2)
            sigma_c = in_c
            image_c = in_c * 4

            # conv weights used for downsampling
            self.weight = torch.zeros(image_c, in_c, 2, 2).to(input)
            for c in range(in_c):
                self.weight[4 * c, c, 0, 0] = 1
                self.weight[4 * c + 1, c, 0, 1] = 1
                self.weight[4 * c + 2, c, 1, 0] = 1
                self.weight[4 * c + 3, c, 1, 1] = 1

            # zeros used to replace repeat
            self.zeros = torch.zeros(in_n, sigma_c, out_h, out_w).to(input)

        self.initialized = True

    def __call__(self, input, sigma):
        assert self.initialized
        output_sigma = self.zeros + sigma
        output_image = torch.nn.functional.conv2d(input, self.weight, stride=2)
        return torch.cat((output_sigma, output_image), dim=1)

---

class Upsample():
    def __init__(self):
        self.initialized = False
        self.weight = None

    def init_weights(self, input):
        with torch.no_grad():
            in_n, in_c, in_h, in_w = input.size()

            image_c = in_c * 4

            self.weight = torch.zeros(in_c + image_c, in_c, 2, 2).to(input)
            for c in range(in_c):
                self.weight[in_c + 4 * c, c, 0, 0] = 1
                self.weight[in_c + 4 * c + 1, c, 0, 1] = 1
                self.weight[in_c + 4 * c + 2, c, 1, 0] = 1
                self.weight[in_c + 4 * c + 3, c, 1, 1] = 1

        self.initialized = True

    def __call__(self, input):
        assert self.initialized
        return torch.nn.functional.conv_transpose2d(input, self.weight, stride=2)

---

在x == upsample(downsample_and_noise_map(x, sigma)) 的意义上，我假设上采样是下采样的倒数（如果我在这个假设中错了，请纠正我）。我还验证了我的下采样版本与你的一致。

# consistency checking code
x = torch.randn(1, 3, 100, 100)
sigma = torch.randn(1)

# OP downsampling
y1 = downsample_and_noise_map(x, sigma)

ds = DownsampleAndNoiseMap()
ds.init_weights(x)
y2 = ds(x, sigma)

print('downsample diff:', torch.sum(torch.abs(y1 - y2)).item())

us = Upsample()
us.init_weights(x)
x_recov = us(ds(x, sigma))

print('recovery error:', torch.sum(torch.abs(x - x_recov)).item())

导致

downsample diff: 0.0
recovery error: 0.0

---

导出到 ONNX

导出时，我们需要在使用torch.onnx.export 之前为新类调用init_weights。例如

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.downsample = DownsampleAndNoiseMap()
        self.upsample = Upsample()
        self.convnet = lambda x: x  # placeholder

    def init_weights(self, x):
        self.downsample.init_weights(x)
        self.upsample.init_weights(x)

    def forward(self, x, sigma):
        x = self.downsample(x, sigma)
        x = self.convnet(x)
        x = self.upsample(x)
        return x

x = torch.randn(1, 3, 100, 100)
sigma = torch.randn(1)

model = Model()
# ... load state dict here
model.init_weights(x)
torch.onnx.export(model, (x, sigma), 'deploy.onnx', verbose=True, input_names=["input", "sigma"], output_names=["output"])

给出 ONNX 图

graph(%input : Float(1, 3, 100, 100)
      %sigma : Float(1)) 
  %2 : Float(1, 3, 50, 50) = onnx::Constant[value=<Tensor>](), scope: Model
  %3 : Float(1, 3, 50, 50) = onnx::Add(%2, %sigma), scope: Model
  %4 : Float(12, 3, 2, 2) = onnx::Constant[value=<Tensor>](), scope: Model
  %5 : Float(1, 12, 50, 50) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[2, 2], pads=[0, 0, 0, 0], strides=[2, 2]](%input, %4), scope: Model
  %6 : Float(1, 15, 50, 50) = onnx::Concat[axis=1](%3, %5), scope: Model
  %7 : Float(15, 3, 2, 2) = onnx::Constant[value=<Tensor>](), scope: Model
  %output : Float(1, 3, 100, 100) = onnx::ConvTranspose[dilations=[1, 1], group=1, kernel_shape=[2, 2], pads=[0, 0, 0, 0], strides=[2, 2]](%6, %7), scope: Model
  return (%output);

---

关于在 iOS 上部署的推荐方式的最后一个问题，我无法回答，因为我没有这方面的经验。

【讨论】：

我能够在 PyTorch 中使用您的代码进行推理，并且还能够成功执行模型转换 PyTorch -> ONNX -> CoreML！耶！有趣的是，我已经尝试使用卷积层来执行下采样。但显然我做错了^^我需要更多时间来详细分析您的解决方案，因为它会引发更多问题。我会尽快回复您。