screamingmamba
screamingmamba

Reputation: 1

Passing input in Tensorflow Mirrored Strategy distirbuted computing

So I am following the example code on tensorflow MirroredStrategy . However, I am getting the following error

raise ValueError('model_fn (%s) must include features argument.' % model_fn) 
ValueError: model_fn (<function build_model_fn_optimizer.<locals>.model_fn at 0x7f1c01098048>) must include features argument.

I think I am following the description however, if the anyone can see any problem in my code snippet and point it out, it will be really helpful.

def NeuralNet(images):
## some neuralnetwork
prediction = 0 # just a place holder it is actually a mask
return prediction


def build_model_fn_optimizer():
    optimizer = tf.train.AdagradOptimizer(learning_rate=0.001)

def model_fn(images, masks, mode):
    predict_mask = NeuralNet(images)

    if mode == tf.estimator.ModeKeys.PREDICT:
        predictions = {'predicted_masks': predict_mask}
        return tf.estimator.EstimatorSpec(mode, predictions = predictions)

    def loss_fn():
        loss = tf.losses.sparse_softmax_cross_entropy(tf.flatten(predict_mask),
                                                      tf.flatten(masks))
        return loss

    if mode == tf.estimator.ModeKeys.EVAL:
        return tf.estimator.EstimatorSpec(mode, loss = loss_fn())

    assert mode == tf.estimator.ModeKeys.TRAIN

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

return model_fn

def main(_):

   distribution = tf.contrib.distribute.MirroredStrategy(
    ["/device:GPU:0", "/device:GPU:1", "/device:GPU:2"])
config = tf.estimator.RunConfig(train_distribute=distribution)

def input_fn():

    base_dir = '/path/to/data'
    images_dir = base_dir + '/images'
    mask_dir = base_dir + '/masks'

    list_fn = os.listdir(images_dir)

    train_fn_list = list_fn[0:5]
    test_fn_list = list_fn[5:7]

    _train_image_fn = []
    _train_mask_fn = []
    for fn in train_fn_list:
        _train_image_fn.append(images_dir + '/' + fn)
        _train_mask_fn.append(mask_dir + '/' + fn)

    _test_image_fn = []
    _test_mask_fn = []
    for fn in test_fn_list:
        _test_image_fn.append(images_dir + '/' + fn)
        _test_mask_fn.append(images_dir + '/' + fn)

    train_images = load_data(_train_image_fn)
    train_masks = load_data(_train_mask_fn)

    images_tf = tf.data.Dataset.from_tensors(images_np)
    masks_tf = tf.data.Dataset.from_tensors(masks_np)
    dataset = tf.data.Dataset.zip((images_tf, masks_tf))
    return dataset

estimator = tf.estimator.Estimator(
    model_fn = build_model_fn_optimizer(), config = config)
estimator.train(input_fn = input_fn, steps =10)


if __name__ == '__main__':
  tf.app.run()

In the function NeuralNet it is actually the network, which returns a mask of the same size of image. I just didnt write the whole network to avoid complications. If someone can help me with this it will be really great. Thanks a bunch.

Upvotes: 0

Views: 1171

Answers (1)

Kevin Yen
Kevin Yen

Reputation: 178

It's not a problem with MirroredStrategy, but rather your model_fn method. It's expected to be

def model_fn(features, labels, mode):

while yours is

def model_fn(images, masks, mode):

hence the error complaining not seeing the argument named features

... must include features argument ...

Upvotes: 3

Related Questions