如何通过Achelous提交一个Tensorflow 任务

Tensorflow 框架简介绍

Tensorflow 是目前人工智能和机器学习领域炙手可热的框架之一。目前其在生物信息领域应用的例子也不断增加。由于配置复杂，部署困难，还涉及到GPU硬件调度细节，一般生物信息学从业人员只能采用单节点形式运行Tensorflow 任务。在单节点模式下，虽然也可以运行Tensorflow任务，但是无法发挥分布式Tensorflow的极限效率。

在Achelous平台上，用户通过partisaner可以轻松进行Tensorflow进行分布式机器学习或者推理。下面通过一个经典的mnist例子演示partisaner的使用。

使用 Partisaner 提交Tensorflow 并行任务

第一步：编写Tensorflow

本示例采用Python 脚本进行编写

部分代码如下：

# encoding:utf-8
import math
import tempfile
import time
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

flags = tf.app.flags
IMAGE_PIXELS = 28
# 定义默认训练参数和数据路径
flags.DEFINE_string('data_dir', '/tmp/mnist', 'Directory  for storing mnist data')
flags.DEFINE_integer('hidden_units', 100, 'Number of units in the hidden layer of the NN')
flags.DEFINE_integer('train_steps', 100000, 'Number of training steps to perform')
flags.DEFINE_integer('batch_size', 100, 'Training batch size ')
flags.DEFINE_float('learning_rate', 0.01, 'Learning rate')
# 定义分布式参数
# 参数服务器parameter server节点
flags.DEFINE_string('ps_hosts', '192.168.32.145:22221', 'Comma-separated list of hostname:port pairs')
# 两个worker节点
flags.DEFINE_string('worker_hosts', '192.168.32.146:22221,192.168.32.160:22221',
                    'Comma-separated list of hostname:port pairs')
# 设置job name参数
flags.DEFINE_string('job_name', None, 'job name: worker or ps')
# 设置任务的索引
flags.DEFINE_integer('task_index', None, 'Index of task within the job')
# 选择异步并行，同步并行
flags.DEFINE_integer("issync", None, "是否采用分布式的同步模式，1表示同步模式，0表示异步模式")
FLAGS = flags.FLAGS
...

为了方便环境变量和参数传递控制，采用shell 脚本进行整合

step=$TRAIN_STEPS
if [ "$1" != "" ]; then
  step=`cat $1`
fi
echo "Start to run tensorflow cluster"
rm -rf /tmp/checkpoint/*
python /tmp/bin/mnist.py --job_name=$TASK_ROLE --task_index=$TASK_INDEX --ps_hosts=$HOSTLIST_PS --worker_hosts=$HOSTLIST_WORKER --train_steps=$step

第二步：编写任务配置文件

{
  "Roles":
  [
    {
      "Name":"ps",
      "Procs":1,
      "Cpu":2,
      "Mem":1000,
      "Server":true
    },
    {
      "Name":"worker",
      "Procs":2,
      "Cpu":1,
      "Mem":1000
    }
  ],
  "Cmd": "sh /tmp/bin/run.sh",
  "Image": "partisaner_demo_tensorflow",
  "User":"root",
  "Env":
  {
    "TRAIN_STEPS":"50000"
  },
  "Volmap":
  {
    "/autofs/vol6/hg-demo/partisaner_demo/tensorflow":"/tmp"
  },
  "AbortOnFail":true,
  "LogDir":"/autofs/vol6/hg-demo/partisaner_demo/tensorflow/log"
}

注意：
Achelous 会自动映射默认环境，用户在无特殊需求的情况下，可以省略环境变量配置工作。

第三步：任务提交

任务提交采用Achelous 中 partisaner进行

mkdir my-tensorflow-demo-job 
cd my-tensorflow-demo-job 
partisaner -c demo-tensorflow-demo.json -t tensorflow

第四步：结果查看

运行结束后，会在程序中指定的输出路径下，生成训练模型文件。结果如下：

>ls -al checkpoint

total 1165
-rw-r--r-- 1 root root    113 May 27  2021 checkpoint
-rw-r--r-- 1 root root  46832 May 27  2021 events.out.tfevents.1622098184.Cc2Apc
-rw-r--r-- 1 root root 136495 May 27  2021 graph.pbtxt
-rw-r--r-- 1 root root 954132 May 27  2021 model.ckpt-0.data-00000-of-00001
-rw-r--r-- 1 root root    515 May 27  2021 model.ckpt-0.index
-rw-r--r-- 1 root root  52500 May 27  2021 model.ckpt-0.meta