モデル並列実行を伴うtf.Estimatorの例

Question

現在、分散型テンソルフローを試しています。 私はtf.contrib.learn.Experimentと一緒にtf.estimator.Estimatorクラス（カスタムモデル関数）を使用しており、それを有効にしてデータ並列実行を実現しました。

しかし、モデルの並列実行を試したいと思います。 Implementation of model parallelism in tensorflowを除いて、私はその事例を見つけることができませんでした。 しかし、tf.estimatorsを使用してこれを実装する方法がわかりません（たとえば、入力機能の処理方法など）。

誰もそれについての経験がありますか、または実例を提供することはできますか？

Answer 1

まずアップcontribがan experimental moduleであり、コアAPIに卒業したクラス（例えばES Estimator）が自動的に非推奨を取得するので、あなたは、tf.estimator.Estimatorの賛成でtf.contrib.learn.Estimatorを使用して停止する必要があります。

あなたの主な質問に戻り、model_fnパラメーターtf.estimator.Estimator.__init__で分散モデルを作成して渡すことができます。

def my_model(features, labels, mode): 
    net = features[X_FEATURE] 
    with tf.device('/device:GPU:1'): 
    for units in [10, 20, 10]: 
     net = tf.layers.dense(net, units=units, activation=tf.nn.relu) 
     net = tf.layers.dropout(net, rate=0.1) 

    with tf.device('/device:GPU:2'): 
    logits = tf.layers.dense(net, 3, activation=None) 
    onehot_labels = tf.one_hot(labels, 3, 1, 0) 
    loss = tf.losses.softmax_cross_entropy(onehot_labels=onehot_labels, 
              logits=logits) 

    optimizer = tf.train.AdagradOptimizer(learning_rate=0.1) 
    train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step()) 
    return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op) 

[...] 

classifier = tf.estimator.Estimator(model_fn=my_model) 

上記モデルは、/device:GPU:2配置と/device:GPU:1配置を6層と3つの他の層を定義します。 my_model関数の戻り値はEstimatorSpecインスタンスにする必要があります。完全な実例はtensorflow examplesにあります。