Передача входных данных в распределенных вычислениях Tensorflow Mirrored Strategy - PullRequest
Новая
       33

Передача входных данных в распределенных вычислениях Tensorflow Mirrored Strategy

0 голосов
/ 25 июня 2018

Итак, я следую примеру кода на tenorflow MirroredStrategy . Тем не менее, я получаю следующую ошибку 

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.

Я думаю, что я следую описанию, однако, если кто-то увидит какую-либо проблему в моем фрагменте кода и укажет на нее, это будет действительно полезно. 

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()

В функции NeuralNet это фактически сеть, которая возвращает маску с таким же размером изображения. Я просто не писал всю сеть, чтобы избежать осложнений. Если кто-то может помочь мне с этим, это будет действительно здорово. Большое спасибо.

1 Ответ

0 голосов
/ 23 января 2019

Это не проблема с MirroredStrategy, а скорее ваш model_fn метод.Ожидается, что он будет 

def model_fn(features, labels, mode):

, в то время как ваш 

def model_fn(images, masks, mode):

, поэтому ошибка, из-за которой не видно аргумента с именем features

... must include features argument ...
...