SageMaker训练MaskRCNN出现TensorFlow Serving容器不可部署告警求解

问题解答 #

你遇到的告警属于该AWS示例训练MaskRCNN模型的正常现象，和你未配置subnets、security_group_ids没有关联：

• 该示例的训练脚本默认仅保存训练过程的checkpoint权重文件，没有按照TensorFlow SavedModel规范导出可部署的模型结构，所以训练结束后SageMaker容器检测到输出目录下的模型不符合TFServing的加载要求，就会抛出该告警，部署失败属于预期结果。
• 你无需重新训练19小时也可以修复该问题，两种可行方案如下：
  1. 直接复用现有训练产物：将训练输出S3路径下的model.tar.gz下载到本地或SageMaker Notebook环境，加载checkpoint权重后，手动调用tf.saved_model.save()接口将模型导出为符合规范的结构：需要在导出根目录下新建版本号目录（比如命名为1），将导出的saved_model.pb和variables子目录放在该版本号目录下，再重新打包为model.tar.gz上传到S3即可用于部署。
  2. 修改训练脚本避免后续训练出现该问题：修改示例中的train.py脚本，在训练结束逻辑后增加SavedModel导出代码，直接将符合结构的模型导出到训练容器的/opt/ml/model路径下，后续训练完成后生成的模型产物就可以直接部署。
• 部署前请确认你使用的SageMaker TensorFlow Serving容器版本和训练时使用的TensorFlow版本完全一致，避免出现兼容性错误。

内容的提问来源于stack exchange，提问作者Xuan Wang