Kubeflow实战系列：利用TFJob导出分布式TensorFlow模型

介绍

本系列将介绍如何在阿里云容器服务上运行Kubeflow, 本文介绍如何使用TfJob导出分布式模型训练模型。

• 第一篇：阿里云上使用JupyterHub
• 第二篇：阿里云上小试TFJob
• 第三篇：利用TFJob运行分布式TensorFlow
• 第四篇：利用TFJob导出分布式TensorFlow模型
• 第五篇：利用TensorFlow Serving进行模型预测

模型训练导出简介

前文中介绍了如何利用TFJob进行分布式的模型训练。对于深度学习的产品化来说，训练只是手段不是目的，目的是将通过训练产生的模型放到手机的程序里或者互联网的应用中，用于语音或者文字的识别等应用场景中。

TensorFlow Serving是Google开源的一个灵活的、高性能的机器学习模型服务系统，能够简化并加速从模型到生产应用的过程。它除了原生支持TensorFlow模型，还可以扩展支持其他类型的机器学习模型。但是TensorFlow Serving支持的模型文件格式是protobuf，而不是TensorFlow模型训练产生的checkpoint文件。这就需要能够将模型训练产生的model.ckpt转化成.pb文件。

在本示例中，我们会提供一个简单的方案示例: 将训练好的checkpoint从用于训练的NAS数据卷迁移到用于模型预测的NAS数据卷，用来给TensorFlow Serving系统加载。这里训练和预测的NAS数据卷，可以选择同一个NAS的不同子目录，只需要通过不同的PV进行标识即可。

创建模型预测使用的NAS

1. 创建NAS数据卷，并且设置与当前Kubernetes集群的同一个具体vpc的挂载点。操作详见文档

2. 在NAS上创建 /serving的数据文件夹, 下载mnist训练所需要的数据

mkdir -p /nfs
mount -t nfs -o vers=4.0 0fc844b526-rqx39.cn-hangzhou.nas.aliyuncs.com:/ /nfs
mkdir -p /nfs/serving
umount /nfs

3. 创建NAS的PV， 以下为示例nas-tf-serving.yaml

apiVersion: v1
kind: PersistentVolume
metadata:
  name: tf-serving-pv
  labels:
    role: tf-serving
spec:
  capacity:
    storage: 10Gi
  accessModes:
    - ReadWriteMany
  storageClassName: nas
  flexVolume:
    driver: "alicloud/nas"
    options:
      mode: "755"
      path: /serving
      server: 0fc844b526-rqx39.cn-hangzhou.nas.aliyuncs.com
      vers: "4.0"

将该模板保存到nas-tf-serving.yaml, 并且创建pv:

# kubectl create -f nas-tf-serving.yaml
persistentvolume "tf-serving-pv" created

4. 利用tf-serving-pvc.yaml创建PVC

kind: PersistentVolumeClaim
apiVersion: v1
metadata:
  name: tf-serving-pvc
spec:
  storageClassName: nas
  accessModes:
    - ReadWriteMany
  resources:
    requests:
      storage: 5Gi
  selector:
    matchLabels:
      role: tf-serving

具体命令：

# kubectl create -f tf-serving-pvc.yaml
persistentvolumeclaim "kubeflow-servin-pvcg" created

查看PVC是否创建成功：

# kubectl get pvc tf-serving-pvc
NAME             STATUS    VOLUME          CAPACITY   ACCESS MODES   STORAGECLASS   AGE
tf-serving-pvc   Bound     tf-serving-pv   10Gi       RWX            nas            28s

5. 创建执行模型导出任务的TFJob

apiVersion: kubeflow.org/v1alpha1
kind: TFJob
metadata:
  name: export-mnist-model
spec:
  replicaSpecs:
    - replicas: 1 # 1 Master
      template:
        spec:
          containers:
            - image: registry.cn-hangzhou.aliyuncs.com/tensorflow-samples/export-mnist-model
              name: tensorflow
              command: ["python", "/app/export_model.py"]
              args:
                - --model_version=1 
                - --checkpoint_path=/training/tensorflow/logs/
                - /serving/mnist
              volumeMounts:
              - name: kubeflow-dist-nas-mnist
                mountPath: "/training"
              - name: tf-serving-pvc
                mountPath: "/serving"
          volumes:
            - name: kubeflow-dist-nas-mnist
              persistentVolumeClaim:
                claimName: kubeflow-dist-nas-mnist
            - name: tf-serving-pvc
              persistentVolumeClaim:
                claimName: tf-serving-pvc
          restartPolicy: Never

• kubeflow-dist-nas-mnist是上一篇分布式训练中保存checkpoint的NAS数据卷，对应的checkpoint文件在其/training/tensorflow/logs/下

6. 查看TFJob

现在可以看到TFJob资源已经被创建了：

# kubectl get tfjob
NAME                 AGE
export-mnist-model   48s

获得该TFJob的RUNTIME ID，这个RUNTIME ID是TFJob和其对应Pod之间的关联

# RUNTIMEID=$(kubectl get tfjob export-mnist-model -o=jsonpath='{.spec.RuntimeId}')

根据RUNTIME ID查询对应执行该训练任务

# kubectl get po -lruntime_id=$RUNTIMEID -a
NAME                                   READY     STATUS    RESTARTS   AGE
export-mnist-model-master-bl1o-0-r7ji0   0/1       Completed   0          4m

在Pod运行过程中，可以通过kubectl logs检查训练模型导出日志, 可以看到模型已经被导出在/serving/mnist/1下:

# kubectl logs export-mnist-model-master-bl1o-0-r7ji0 
/usr/local/lib/python2.7/dist-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.
  from ._conv import register_converters as _register_converters
2018-06-13 13:43:45.183857: I tensorflow/core/platform/cpu_feature_guard.cc:137] Your CPU supports instructions that this TensorFlow binary was not compiled to use: SSE4.1 SSE4.2 AVX AVX2 FMA
WARNING:tensorflow:From /app/export_model.py:74: initialize_all_tables (from tensorflow.python.ops.lookup_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.tables_initializer` instead.
Tensor("input/x-input:0", shape=(?, 784), dtype=float32)
Tensor("cross_entropy/logits:0", shape=(?, 10), dtype=float32)
Exporting trained model to /serving/mnist/1
Done exporting!

7. 查看模型

登录到NAS上检查模型文件

# tree serving
serving
└── mnist
    └── 1
        ├── saved_model.pb
        └── variables
            ├── variables.data-00000-of-00001
            └── variables.index

并且可以通过saved_model_cli检查模型定义, 这里可以看到输入的input是images,输出的结果是scores。这个在模型预测中，写grpc客户端会用到

# saved_model_cli show --dir /serving/mnist/1 --all

MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:

signature_def['predict_images']:
The given SavedModel SignatureDef contains the following input(s):
inputs['images'] tensor_info:
    dtype: DT_FLOAT
    shape: (-1, 784)
    name: input/x-input:0
The given SavedModel SignatureDef contains the following output(s):
outputs['scores'] tensor_info:
    dtype: DT_FLOAT
    shape: (-1, 10)
    name: cross_entropy/logits:0
Method name is: tensorflow/serving/predict

总结

模型导出是从训练转向服务的桥梁，通过使用TFJob完成分布式模型训练的导出，可以标准化和自动化这部分工作，从而有机会建立从模型训练，验证，导出，预测的工作流。在后面的文章中，我们会介绍如何通过TensorFlow Serving使用导出的模型进行预测。