如何将已保存的模型从NCHW转换为NHWC？

Question

我在GPU上训练了NCHW模型并保存了最佳状态。 现在，我想在CPU上进行推断，该CPU显然仅支持NHWC（提及该错误）。 我必须使用NHWC重新训练模型，还是可以转换模型？

Answer 1

我处在相同的情况下，在只有CPU可用的实例上尝试在经过GPU训练的模型上运行model.predict时会看到类似的错误：

tensorflow.python.framework.errors_impl.InvalidArgumentError: Default MaxPoolingOp only supports NHWC on device type CPU

我最终发现，在Intel CPU上，只要启用了MKL，就可以成功地将模型应用于NCHW格式的数据。使用pip，可以使用以下命令安装启用MKL的tensorflow：

pip install intel-tensorflow

您可以使用以下方法检查它是否已启用（在tensorflow 2.3中）：

tf.python._pywrap_util_port.IsMklEnabled()