我在MirroredStrategy中使用GANEstimator在单个实例的多个GPU上工作。在我的情况下，input_fn是具有以下设置的tf.data.Dataset:
dataset = dataset.repeat()
dataset = dataset.shuffle(buffer_size=100)
dataset = dataset.batch(self.batch_size, drop_remainder=True)
dataset = dataset.prefetch(100)
我之所以这样问，是因为我需要手动指定诸如
dataset.shard()之类的东西才能将不同的数据传递给工作人员吗？我正在研究Estimator和MirroredStrategy的代码，但是我不清楚发生了什么。从description of distributed strategies还会造成其他困惑:MirroredStrategy: This does in-graph replication with synchronous
training on many GPUs on one machine. Essentially, we create copies of all
variables in the model's layers on each device. We then use all-reduce
to combine gradients across the devices before applying them
to the variables to keep them in sync.
CollectiveAllReduceStrategy: This is a version of MirroredStrategy
for multi-worker training.
那么MirroredStratedy是否只使用一名 worker ？我不明白我需要指定等于一个塔的容量的批处理大小，否则得到OOM。有人可以给我指出代码，并解释一下这种简单的设置如何与批处理一起工作:
def create_dataset():
...
dataset = dataset.repeat()
dataset = dataset.shuffle(buffer_size=100)
dataset = dataset.batch(self.batch_size, drop_remainder=True)
dataset = dataset.prefetch(100)
return dataset
NUM_GPUS = 4
strategy = tf.contrib.distribute.MirroredStrategy(num_gpus=NUM_GPUS)
optimizer = tf.train.RMSPropOptimizer(learning_rate=0.01, use_locking=True)
optimizer_d = tf.train.RMSPropOptimizer(learning_rate=0.01, use_locking=True)
config = tf.estimator.RunConfig(save_checkpoints_steps=100,
save_summary_steps=1, keep_checkpoint_max=50,
train_distribute=strategy)
# I have more hooks here, just simplified to show
def get_hooks_fn(GANTrainOps):
disjoint_train_hook_func = tfgan.get_sequential_train_hooks(
train_steps=tfgan.GANTrainSteps(10, 1)
) # g steps, d steps
disjoint_train_hooks = disjoint_train_hook_func(GANTrainOps)
return [update_hook, summary_hook] + disjoint_train_hooks
# Create GAN estimator.
gan_estimator = tfgan.estimator.GANEstimator(
model_dir = '/data/checkpoints/estimator_model',
generator_fn = generator_fn,
discriminator_fn = discriminator_fn,
generator_loss_fn = generator_loss_fn,
discriminator_loss_fn = discriminator_loss_fn,
generator_optimizer = optimizer,
discriminator_optimizer = optimizer_d,
use_loss_summaries=True,
config=config,
get_hooks_fn=get_hooks_fn)
gan_estimator.train(input_fn=create_dataset, steps=10000)
谢谢!
MirroredStrategy的代码包含:
1)怪异的措辞:
2)
默认情况下，此参数为False。
编辑:
到目前为止，我发现一段时间后
tf.estimator.train() 指向似乎是strategy.make_input_fn_iterator()的内容:def _get_iterator_from_input_fn(self, input_fn, mode, distribution=None):
if distribution is not None:
iterator = distribution.make_input_fn_iterator(
lambda _: self._call_input_fn(input_fn, mode))
input_hooks = [
estimator_util.DistributedIteratorInitializerHook(iterator)]
else:
result = self._call_input_fn(input_fn, mode)
iterator = result.make_initializable_iterator()
input_hooks = [estimator_util._DatasetInitializerHook(iterator)]
return iterator, input_hooks
make_input_fn_iterator()但是它已从MirroredStrategy的代码中删除，并且不再存在!我不明白它是如何工作的以及数据集实际上是在哪里拆分的。
EDIT2:我无法使用grep在tensorflow 1.12.0发行版中找到
make_input_fn_iterator行。似乎代码中完全没有它。 最佳答案
好的，花了一些时间研究github后，我发现它已经与我的tf 1.12.0不同。因此，进入1.12.0的本地文件可以得到:
GANEstimator继承了tf.python.estimator.Estimator
Estimator.init():
# The distribute field contains an instance of DistributionStrategy.
self._train_distribution = self._config.train_distribute
那么向下的路径是:
tf.contrib.gan.GANEstimator -> tf.python.estimator.Estimator.train() -->
tf.python.estimator.Estimator._train_model(input_fn, hooks, saving_listeners) -->
._train_model_distributed(input_fn, hooks, saving_listeners) -->
._get_iterator_from_input_fn(input_fn, model_fn_lib.ModeKeys.TRAIN, self._train_distribution) -->
distribution.distribute_dataset(lambda: self._call_input_fn(input_fn, mode))
在我的情况下需要
MirrorredStrategy.distribute_dataset():def distribute_dataset(self, dataset_fn):
if self._cluster_spec:
return values.MultiWorkerDataset(
partial(self._call_dataset_fn, dataset_fn), self._worker_device_map,
self._prefetch_on_device, self._auto_shard_dataset)
else:
return values.PerDeviceDataset(
self._call_dataset_fn(dataset_fn), self._devices,
self._prefetch_on_device)
tensorflow/python/training/distribute.py: def _call_dataset_fn(self, dataset_fn):
result = dataset_fn()
if not isinstance(result, dataset_ops.Dataset):
raise ValueError(
"dataset_fn() must return a tf.data.Dataset when using a "
"DistributionStrategy.")
return result
我假设使用
PerDeviceDataset，所以最后我在values.py中找到了这两个类:class PerDeviceDataset(object):
"""Like `tf.data.Dataset` split devices, producing `PerDevice` data."""
def __init__(self, dataset, devices, prefetch_on_device=None):
self._devices = devices
# Default to using prefetching in graph mode, unless specified.
# TODO(priyag): Enable prefetching in eager mode.
self._prefetch_on_device = prefetch_on_device
if self._prefetch_on_device is None:
self._prefetch_on_device = not context.executing_eagerly()
assert not (self._prefetch_on_device and context.executing_eagerly()), (
"Prefetching is only supported in graph mode currently")
if self._prefetch_on_device:
self._dataset = dataset.apply(
prefetching_ops_v2.prefetch_to_devices(self._devices))
else:
# TODO(priyag): If dropping remainder is not appropriate, find another
# approach to distributing the dataset when not possible to divide evenly.
# Possibly not an issue when we start using PartitionedDataset.
self._dataset = dataset.batch(len(devices), drop_remainder=True)
def make_one_shot_iterator(self):
"""Get a one time use iterator for the distributed PerDeviceDataset."""
dataset_iterator = self._dataset.make_one_shot_iterator()
return PerDeviceDataIterator(dataset_iterator, self._devices,
self._prefetch_on_device)
def make_initializable_iterator(self):
"""Get an initializable iterator for the distributed PerDeviceDataset."""
dataset_iterator = self._dataset.make_initializable_iterator()
return PerDeviceDataIterator(dataset_iterator, self._devices,
self._prefetch_on_device)
class PerDeviceDataIterator(object):
"""An iterator (like `tf.data.Iterator`) into a `PerDeviceDataset`."""
def __init__(self, iterator, devices, prefetch_on_device=None):
self._iterator = iterator
self._devices = devices
self._prefetch_on_device = prefetch_on_device
@property
def initializer(self):
return self._iterator.initializer
def get_next(self, name=None):
"""Scatter the input across devices."""
if self._prefetch_on_device:
data_list = self._iterator.get_next(name=name)
index = dict(zip(self._devices, data_list))
else:
batch = self._iterator.get_next(name=name)
index = {}
def get_ith(i):
return lambda x: x[i]
for i, d in enumerate(self._devices):
index[d] = nest.map_structure(get_ith(i), batch)
if context.executing_eagerly():
with ops.device(d):
index[d] = nest.map_structure(array_ops.identity, index[d])
return regroup(index)
因此，据我了解，首先，我的
dataset_fn()函数仅被调用来获取数据集对象，然后在其之上应用一批具有GPU数量的大小。该批次的元素必须是我在dataset_fn()内的数据集初始化中定义的实际批次，已分配给不同的设备。关于tensorflow - 使用MirroredStrategy时，tensorflow Estimator是否为工作人员分配了不同的批次？，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/54327610/