有什么方法可以将 .pb 文件的数据格式从 NCHW 转换为 NHWC？

Question

我有一个基于数据格式N(batch) x C(channel) x H(height) x W(width) 在 Pytorch 中训练的 CNN 模型。我将预训练模型保存为 model.pth。之后，我使用现有函数从 model.pth -> model.onnx 转换了预训练模型：

torch.onnx.export(model, dummy_input, "model.onnx")

然后，我通过下面的模块转换了这个 model.onnx -> model.pb：

import onnx
from onnx_tf.backend import prepare 

model_onnx = onnx.load('model.onnx')
tf_rep = prepare(model_onnx)
tf_rep.export_graph('model.pb')

问题是：我想在 CPU 设备上使用这个 model.pb，它需要 NHWC 数据格式。但是，我的模型是基于 NCHW 数据格式的。有没有什么方法可以把这个model.pb的数据格式从NCHW转换成NHWC？

Answer 1

简短的回答，你处境艰难。

答案很长，虽然很难，但可能。使您的问题变得困难的是您的图表已经过训练。在创建训练图时转换 NCHW->NHWC 效率低，但更容易。查看类似的答案here 和here。

现在回答您的问题，您必须使用自定义卷积运算符重载 conv2D 运算符。这是一个开始使用的伪代码。

   tensor Conv2D(X, W, B) {
     int perm[] = {0, 3, 1, 2};
     X = transposeTensor(X, perm);
     W = transposeTensor(W, perm);
     Y = Conv2D_orig(X, W, B, ...) ;
     perm = {0, 2, 3, 1};
     return transposeTensor(Y, perm);
   }

Answer 2

你可以只转置输入张量吗？例如input.transpose(1,2).transpose(2,3)?

>>> torch.randn( (3,3,3,3), names=['n','c','h','w']).transpose(1,2).transpose(2,3).names
('n', 'h', 'w', 'c')