尝试在tf.estimator中使用MirroredStrategy时出错

时间:2018-04-12 21:13:02

标签: python tensorflow tensorflow-estimator

我正在尝试使用tf.contrib.distribute.MirroredStrategy作为tf.estimator.RunConfig的参数,为我的tensorflow训练代码添加multi-gpu支持。

Tensorflow版本:1.7(从源代码编译)

Python版本:3.5

OS平台和版本:Linux Ubuntu 16.04.2

我收到以下错误消息:

Traceback (most recent call last):
  File "python3.5/site-packages/tensorflow/python/training/coordinator.py", line 297, in stop_on_exception
    yield
  File "python3.5/site-packages/tensorflow/contrib/distribute/python/mirrored_strategy.py", line 248, in _call_for_each_tower
    self, *merge_args, **merge_kwargs)
  File "python3.5/site-packages/tensorflow/python/training/optimizer.py", line 667, in _distributed_apply
    reduced_grads = distribution.batch_reduce("sum", grads_and_vars)
  File "python3.5/site-packages/tensorflow/python/training/distribute.py", line 801, in batch_reduce
    return self._batch_reduce(method_string, value_destination_pairs)
  File "python3.5/site-packages/tensorflow/contrib/distribute/python/mirrored_strategy.py", line 295, in _batch_reduce
    value_destination_pairs)
  File "python3.5/site-packages/tensorflow/contrib/distribute/python/cross_tower_ops.py", line 169, in batch_reduce
    raise ValueError("`value_destination_pairs` must be a list or a tuple of "
ValueError: `value_destination_pairs` must be a list or a tuple of tuples of PerDevice objects and destinations

以下代码产生错误(我省略了将tfrecord解析为图像张量的代码,因为我不相信此代码会影响错误,但我可以在必要时添加它):

import glob, os
import tensorflow as tf
slim = tf.contrib.slim

# ...
# definition of args (arguments parser)

def input_fn():

    dataset = tf.data.TFRecordDataset(glob.glob(os.path.join(args.train_data_dir, 'train*')))
    dataset = dataset.map(
                lambda x: parse_and_preprocess_image(x, args.image_size),
                num_parallel_calls=2,
    )
    dataset = dataset.repeat()
    dataset = dataset.batch(batch_size=4)
    dataset = dataset.prefetch(1)

    return dataset


def model_fn(features, labels=None, mode=tf.estimator.ModeKeys.TRAIN, params=None):

    train_images_batch = features
    res = slim.conv2d(inputs=train_images_batch, kernel_size=9, stride=1, num_outputs=3, scope='conv1')
    loss = tf.reduce_mean((train_images_batch - res) ** 2)
    optimizer = tf.train.AdamOptimizer(0.001)
    train_op = slim.learning.create_train_op(loss, optimizer)
    return tf.estimator.EstimatorSpec(
        mode=tf.estimator.ModeKeys.TRAIN,
        loss=loss, train_op=train_op)


def train():

    init()

    distribution = tf.contrib.distribute.MirroredStrategy(num_gpus=args.num_gpus)

    config = tf.estimator.RunConfig(
        model_dir=args.log_dir,
        train_distribute=distribution,
    )

    estimator = tf.estimator.Estimator(model_fn=model_fn, config=config)
    estimator.train(
            input_fn=input_fn,
            max_steps=args.train_steps,
        )


def main():
    add_arguments()
    train()


if __name__ == '__main__':
    main()

谢谢!

ADVA

1 个答案:

答案 0 :(得分:1)

如果您指定num_gpus=1,则会发生此错误。对于单个GPU,可以使用OneDeviceStrategy("/device:GPU:0")代替MirroredStrategy

相关问题