TensorFlow官方ResNet模型在导出saved_model时的注意事项。

导出命令示例:

python3 posenet_main.py \
  --data_dir=/data/dataset/public/facial_landmark/tfrecord/ \
  --model_dir=/data/posenet/posenet_train/ \
  --train_epochs=0 \
  --export_dir=/data/posenet/exported/ \
  --image_bytes_as_serving_input=True \
  --resnet_size=18 \
  --resnet_version=2 \
  --data_format=channels_last

使用未解码的图像文件作为输入

指定 --image_bytes_as_serving_input=True

自定义图像预处理

如果你的网络需要图像输入前做预处理的话,可以同时修改 resnet_run_loop.py 这个模块中的image_bytes_serving_input_fn 函数。

导出在CPU端推演的模型

指定 --data_format=channels_last

官方的ResNet模型同时支持NHWC与HWCN。第一种在GPU上更快,是在CUDA环境下导出模型的默认选择。如果你在GPU下训练模型,又要在CPU下进行推演,记得指定这个参数。否则会有类似下方的错误:

E external/org_tensorflow/tensorflow/core/common_runtime/executor.cc:624] Executor failed to create kernel. Invalid argument: Default MaxPoolingOp only supports NHWC on device type CPU
	 [[{{node resnet_model/max_pooling2d/MaxPool}}]]

Reference

封图:Photo by Leonard von Bibra on Unsplash