这将是一篇很长的帖子，在此先对不起...
我正在研究一种去噪算法，我的目标是：
去噪算法包括以下三个部分：
第一部分的想法很简单，但解释起来并不容易。例如，给定输入彩色图像和代表图像噪声的标准偏差的输入值“ 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标准：批处理，通道，高度，宽度）和sigma值为0.1的输入张量。输出张量具有以下属性：
如果此代码不够清晰，我可以发布一个更幼稚的版本。
上采样部分是下采样部分的倒数功能。
我能够在PyTorch中使用此功能进行培训和测试。
然后，我尝试使用ONNX将模型转换为CoreML，作为中间步骤。 转换为ONNX会生成“ TracerWarning”。从ONNX到CoreML的转换失败（TypeError：1.0的类型为numpy.float64，但应为以下类型之一：int，long）。问题出自下采样+噪声水平图（以及上采样）。
当我删除下采样+噪声电平图和上采样层时，由于仅保留了简单的卷积网络，因此能够非常轻松地转换为ONNX和CoreML。这意味着我有一个解决问题的方法：在移动端使用2个着色器实现这2层。但是我对这个解决方案不满意，因为我希望我的模型包含所有图层^^
在考虑在此处撰写帖子之前，我已经在Internet上爬网以找到答案，并且能够使用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方法是什么？感谢您的帮助（和耐心等待）^^
答案 0 :(得分：1)
免责声明我不熟悉CoreML或部署到iOS，但是我确实有通过ONNX在TensorRT和OpenVINO中部署PyTorch模型的经验。
在部署到其他框架时，我面临的主要问题是切片和重复张量之类的操作在其他框架中的支持有限。通常，我们可以构造等效的conv或transpose-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
导出时，我们需要在使用init_weights之前为新类调用torch.onnx.export。例如
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上部署的方式的最后一个问题，由于我在该领域没有经验，所以我无法回答。