tf.contrib.learn 分布式模型训练任务快速入门

一个简单例子

tf.contrib.learn 的模型训练任务在执行训练的时候需要调用 Estimator 的 fit 方法，调用该方法时指定一个 input_fn 来指定模型的输入（features）和预期输出（labels）。

import tensorflow as tf
import numpy as np
tf.logging.set_verbosity(tf.logging.INFO)

features = [tf.contrib.layers.real_valued_column("x", dimension=1)]
estimator = tf.contrib.learn.LinearRegressor(feature_columns=features)

x = np.array([1., 2., 3., 4.])
y = np.array([0., -1., -2., -3.])
input_fn = tf.contrib.learn.io.numpy_input_fn({"x":x}, y, batch_size=4,
                                              num_epochs=1000)
estimator.fit(input_fn=input_fn, steps=1000)
estimator.evaluate(input_fn=input_fn)

在编写支持 TaaS 平台的 tf.contrib.learn 模型训练任务的时候，只需要将传递给 Estimator 的 fit 和 evaluate 方法的 input_fn 参数来定义一个 dist_base.Experiment 对象即可。

import tensorflow as tf
import numpy as np
from caicloud.clever.tensorflow import dist_base
tf.logging.set_verbosity(tf.logging.INFO)

features = [tf.contrib.layers.real_valued_column("x", dimension=1)]
estimator = tf.contrib.learn.LinearRegressor(feature_columns=features)

x = np.array([1., 2., 3., 4.])
y = np.array([0., -1., -2., -3.])
input_fn = tf.contrib.learn.io.numpy_input_fn({"x":x}, y, batch_size=4,
                                              num_epochs=1000)
experiment = dist_base.Experiment(
    estimator = estimator,
    train_input_fn = input_fn,
    eval_input_fn = input_fn)
experiment.run()

Boston House 样例

TensorFlow 官方文档提供 Boston House Values 样例来讲解了 tf.contrib.learn 的使用（请参数 Building Input Functions with tf.contrib.learn），实现代码请参考 boston.py。

下面通过改造 Boston House 样例代码来支持在 TaaS 平台上执行。

# coding=utf-8
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import itertools

import pandas as pd
import tensorflow as tf
import tensorflow.contrib.learn as learn

from caicloud.clever.tensorflow import dist_base

tf.app.flags.DEFINE_string("data_dir",
                           ".",
                           "data directory path.")
FLAGS = tf.app.flags.FLAGS

tf.logging.set_verbosity(tf.logging.INFO)

COLUMNS = ["crim", "zn", "indus", "nox", "rm", "age",
           "dis", "tax", "ptratio", "medv"]
FEATURES = ["crim", "zn", "indus", "nox", "rm",
            "age", "dis", "tax", "ptratio"]
LABEL = "medv"

# 加载训练数据集和测试数据集
training_set = pd.read_csv("{0}/boston_train.csv".format(FLAGS.data_dir),
                           skipinitialspace=True,
                           skiprows=1, names=COLUMNS)
test_set = pd.read_csv("{0}/boston_test.csv".format(FLAGS.data_dir),
                       skipinitialspace=True,
                       skiprows=1, names=COLUMNS)

# 输入数据的特征列表
feature_cols = [tf.contrib.layers.real_valued_column(k)
                for k in FEATURES]

# 定义深度神经网络回归模型
run_config = tf.contrib.learn.RunConfig(
    save_checkpoints_secs=dist_base.cfg.save_checkpoints_secs)
regressor = tf.contrib.learn.DNNRegressor(
    feature_columns=feature_cols,
    hidden_units=[10, 10],
    model_dir=dist_base.cfg.logdir,
    config=run_config)

def input_fn(data_set):
    feature_cols = {k: tf.constant(data_set[k].values)
                    for k in FEATURES}
    labels = tf.constant(data_set[LABEL].values)
    return feature_cols, labels

# 定义 dist_base.Experiment 对象来分布式执行 tf.contrib.learn 模型训练任务。
# 指定相对应的 esimator 对象，将要传递给 estimator.fit() 方法进行模型训练的
# input_fn 参数传递给 dist_base.Experiment 的 train_input_fn 参数。将执行模型
# 评估时传递给 estimator.evaluate() 方法的 input_fn 参数传递给 dist_base.Experiment
# 对象的 eval_input_fn 参数。
exp = dist_base.Experiment(
    estimator = regressor,
    train_input_fn = lambda: input_fn(training_set),
    eval_input_fn = lambda: input_fn(test_set),
    eval_steps = 1)

# 调用 dist.Experiment.run() 方法来分布式执行 tf.contrib.learn 模型训练任务。
exp.run()