blog-distributed-tensorflow-map-4.py

with tf.device(tf.train.replica_device_setter(
                    worker_device="/job:worker/task:%d" % task_index,
                    cluster=cluster)):

    def build_model():
        model_input = tf.placeholder(tf.float32,\
                    [None,args['num_features'] ])
        model_labels = tf.placeholder(tf.float32, [None, args['num_classes']    ])
        logits = tf.keras.layers.Dense(args['num_classes'])(model_input)
        model_output = tf.nn.softmax(logits)

        tf_global_step = tf.train.get_or_create_global_step()
        tf_loss = tf.reduce_mean(tf.nn.weighted_cross_entropy_with_logits(logits=logits, targets=model_labels,pos_weight=weights))
        tf_optimizer = tf.train.AdamOptimizer(learning_rate=args['learning_rate']).minimize(tf_loss,
                                                                            global_step=tf.train.get_global_step())

    model_input,\
    model_labels,\
    model_output,\
    tf_global_step,\
    tf_loss,\
    tf_optimizer,\
    tf_metrics = build_model()