如何将python中的NHWC转换为NCHW？

Question

我有一个TensorFlow Keras模型，它以.pb格式存储，从.pb格式，我使用tf2onnx模型将模型转换成.onnx格式

!python -m tf2onnx.convert --saved-model model.pb --output model.onnx 

现在，在转换之后，我看到我的输入层是NHWC格式的，我需要将它转换成NCHW，以实现我正在使用的

!python -m tf2onnx.convert --saved-model model.pb --output model_3.onnx --inputs-as-nchw input0:0

这仍然给我相同的输出，与NHWC一样，我不得不在中使用上面的模型，它只接受NCHW格式。

我发现这个链接讨论了输入层的转置，但不幸的是，这也不起作用。Convert between NHWC and NCHW in TensorFlow

#import tensorflow as tf
images_nhwc = tf.compat.v1.placeholder(tf.float32, [1, 200, 300, 3])  
# input batch
out = tf.transpose(images_nhwc, [0, 3, 1, 2])
#print(out.get_shape())
model.build(out.get_shape())

​

如果一些专家能够分享他们对如何将NHWC转化为NCHW的想法，那将是非常有帮助的。

Answer 1

我找到了解决办法。我不得不使用tf2onnx.convert.from_keras的最新代码。我从tf2onnx拿走了主干道

!pip install --force-reinstall  git+https://github.com/onnx/tensorflow-onnx.git@main
!pip show tf2onnx
!pip freeze | grep tf2onnx

一旦完成，我就能够在https://github.com/onnx/tensorflow-onnx/tree/e896723e410a59a600d1a73657f9965a3cbf2c3b上加载最新的功能和更新的代码。

下面是用于将模型从.pb转换为.onnx的代码，以及将NHWC转换为NCHW的代码。

# give the list of *inputs* which should be converted and returned *as nchw*
_INPUT = model.input.name

model_proto, external_tensor_storage = tf2onnx.convert.from_keras(model, inputs_as_nchw=[_INPUT])

上面的代码最大的问题是_INPUT，它应该是一个列表，我能够在测试用例中找到这些信息。