distribution_utils.py 12.8 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Helper functions for running models in a distributed setting."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

21
22
import json
import os
23
24
import random
import string
25
import tensorflow as tf
26
from tensorflow.contrib import distribute as contrib_distribute
27

28
29
from official.utils.misc import tpu_lib

30

31
32
def _collective_communication(all_reduce_alg):
  """Return a CollectiveCommunication based on all_reduce_alg.
33

34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
  Args:
    all_reduce_alg: a string specifying which collective communication to pick,
      or None.

  Returns:
    tf.distribute.experimental.CollectiveCommunication object

  Raises:
    ValueError: if `all_reduce_alg` not in [None, 'ring', 'nccl']
  """
  collective_communication_options = {
      None: tf.distribute.experimental.CollectiveCommunication.AUTO,
      "ring": tf.distribute.experimental.CollectiveCommunication.RING,
      "nccl": tf.distribute.experimental.CollectiveCommunication.NCCL
  }
  if all_reduce_alg not in collective_communication_options:
    raise ValueError(
        "When used with `multi_worker_mirrored`, valid values for "
        "all_reduce_alg are ['ring', 'nccl'].  Supplied value: {}".format(
            all_reduce_alg))
  return collective_communication_options[all_reduce_alg]


def _mirrored_cross_device_ops(all_reduce_alg, num_packs):
  """Return a CrossDeviceOps based on all_reduce_alg and num_packs.

  Args:
    all_reduce_alg: a string specifying which cross device op to pick, or None.
    num_packs: an integer specifying number of packs for the cross device op.

  Returns:
    tf.distribute.CrossDeviceOps object or None.

  Raises:
    ValueError: if `all_reduce_alg` not in [None, 'nccl', 'hierarchical_copy'].
  """
  if all_reduce_alg is None:
    return None
  mirrored_all_reduce_options = {
      "nccl": tf.distribute.NcclAllReduce,
      "hierarchical_copy": tf.distribute.HierarchicalCopyAllReduce
  }
  if all_reduce_alg not in mirrored_all_reduce_options:
    raise ValueError(
        "When used with `mirrored`, valid values for all_reduce_alg are "
        "['nccl', 'hierarchical_copy'].  Supplied value: {}".format(
            all_reduce_alg))
  cross_device_ops_class = mirrored_all_reduce_options[all_reduce_alg]
  return cross_device_ops_class(num_packs=num_packs)
83

84

85
def get_distribution_strategy(distribution_strategy="mirrored",
86
                              num_gpus=0,
87
                              num_workers=1,
88
                              all_reduce_alg=None,
89
90
                              num_packs=1,
                              tpu_address=None):
91
92
93
  """Return a DistributionStrategy for running the model.

  Args:
94
    distribution_strategy: a string specifying which distribution strategy to
95
      use. Accepted values are 'off', 'one_device', 'mirrored',
96
      'parameter_server', 'multi_worker_mirrored', and 'tpu' -- case insensitive.
97
      'off' means not to use Distribution Strategy; 'tpu' means to use
98
      TPUStrategy using `tpu_address`.
99
    num_gpus: Number of GPUs to run this model.
100
    num_workers: Number of workers to run this model.
101
102
103
104
105
    all_reduce_alg: Optional. Specifies which algorithm to use when performing
      all-reduce. For `MirroredStrategy`, valid values are "nccl" and
      "hierarchical_copy". For `MultiWorkerMirroredStrategy`, valid values are
      "ring" and "nccl".  If None, DistributionStrategy will choose based on
      device topology.
106
107
    num_packs: Optional.  Sets the `num_packs` in `tf.distribute.NcclAllReduce`
      or `tf.distribute.HierarchicalCopyAllReduce` for `MirroredStrategy`.
108
109
    tpu_address: Optional. String that represents TPU to connect to. Must not
      be None if `distribution_strategy` is set to `tpu`.
110
  Returns:
111
    tf.distribute.DistibutionStrategy object.
Shining Sun's avatar
Shining Sun committed
112
  Raises:
113
    ValueError: if `distribution_strategy` is 'off' or 'one_device' and
114
115
      `num_gpus` is larger than 1; or `num_gpus` is negative or if
      `distribution_strategy` is `tpu` but `tpu_address` is not specified.
116
  """
117
118
119
120
121
  if num_gpus < 0:
    raise ValueError("`num_gpus` can not be negative.")

  distribution_strategy = distribution_strategy.lower()
  if distribution_strategy == "off":
122
    if num_gpus > 1:
123
124
125
      raise ValueError(
          "When {} GPUs and  {} workers are specified, distribution_strategy "
          "flag cannot be set to 'off'.".format(num_gpus, num_workers))
126
127
    return None

128
  if distribution_strategy == "tpu":
Hongkun Yu's avatar
Hongkun Yu committed
129
    # When tpu_address is an empty string, we communicate with local TPUs.
130
131
132
    cluster_resolver = tpu_lib.tpu_initialize(tpu_address)
    return tf.distribute.experimental.TPUStrategy(cluster_resolver)

133
  if distribution_strategy == "multi_worker_mirrored":
134
    return tf.distribute.experimental.MultiWorkerMirroredStrategy(
135
        communication=_collective_communication(all_reduce_alg))
136

137
  if distribution_strategy == "one_device":
138
    if num_gpus == 0:
Toby Boyd's avatar
Toby Boyd committed
139
      return tf.distribute.OneDeviceStrategy("device:CPU:0")
140
141
142
143
    if num_gpus > 1:
      raise ValueError("`OneDeviceStrategy` can not be used for more than "
                       "one device.")
    return tf.distribute.OneDeviceStrategy("device:GPU:0")
144

145
  if distribution_strategy == "mirrored":
146
147
    if num_gpus == 0:
      devices = ["device:CPU:0"]
Shining Sun's avatar
Shining Sun committed
148
    else:
149
      devices = ["device:GPU:%d" % i for i in range(num_gpus)]
150
151
    return tf.distribute.MirroredStrategy(
        devices=devices,
152
        cross_device_ops=_mirrored_cross_device_ops(all_reduce_alg, num_packs))
153

154
  if distribution_strategy == "parameter_server":
155
    return tf.distribute.experimental.ParameterServerStrategy()
156
157
158
159

  raise ValueError(
      "Unrecognized Distribution Strategy: %r" % distribution_strategy)

160

161
def per_replica_batch_size(batch_size, num_gpus):
162
163
  """For multi-gpu, batch-size must be a multiple of the number of GPUs.

164
165
166

  Note that distribution strategy handles this automatically when used with
  Keras. For using with Estimator, we need to get per GPU batch.
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183

  Args:
    batch_size: Global batch size to be divided among devices. This should be
      equal to num_gpus times the single-GPU batch_size for multi-gpu training.
    num_gpus: How many GPUs are used with DistributionStrategies.

  Returns:
    Batch size per device.

  Raises:
    ValueError: if batch_size is not divisible by number of devices
  """
  if num_gpus <= 1:
    return batch_size

  remainder = batch_size % num_gpus
  if remainder:
Toby Boyd's avatar
Toby Boyd committed
184
185
186
    err = ('When running with multiple GPUs, batch size '
           'must be a multiple of the number of available GPUs. Found {} '
           'GPUs with a batch size of {}; try --batch_size={} instead.'
187
188
189
          ).format(num_gpus, batch_size, batch_size - remainder)
    raise ValueError(err)
  return int(batch_size / num_gpus)
190

Toby Boyd's avatar
Toby Boyd committed
191

192
193
194
195
196
197
198
199
200
# The `SyntheticDataset` is a temporary solution for generating synthetic data
# directly on devices. It is only useful for Keras with Distribution
# Strategies. We will have better support in `tf.data` or Distribution Strategy
# later.
class SyntheticDataset(object):
  """A dataset that generates synthetic data on each device."""

  def __init__(self, dataset, split_by=1):
    # dataset.take(1) doesn't have GPU kernel.
Toby Boyd's avatar
Toby Boyd committed
201
    with tf.device('device:CPU:0'):
202
203
204
      tensor = tf.data.experimental.get_single_element(dataset.take(1))
    flat_tensor = tf.nest.flatten(tensor)
    variable_data = []
205
    initializers = []
206
207
208
    for t in flat_tensor:
      rebatched_t = tf.split(t, num_or_size_splits=split_by, axis=0)[0]
      assert rebatched_t.shape.is_fully_defined(), rebatched_t.shape
209
      v = tf.compat.v1.get_local_variable(self._random_name(),
Toby Boyd's avatar
Toby Boyd committed
210
                                          initializer=rebatched_t)
211
      variable_data.append(v)
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
      initializers.append(v.initializer)
    input_data = tf.nest.pack_sequence_as(tensor, variable_data)
    self._iterator = SyntheticIterator(input_data, initializers)

  def _random_name(self, size=10, chars=string.ascii_uppercase + string.digits):
    return ''.join(random.choice(chars) for _ in range(size))

  def __iter__(self):
    return self._iterator

  def make_one_shot_iterator(self):
    return self._iterator

  def make_initializable_iterator(self):
    return self._iterator


class SyntheticIterator(object):
  """A dataset that generates synthetic data on each device."""

  def __init__(self, input_data, initializers):
    self._input_data = input_data
    self._initializers = initializers
235
236
237
238

  def get_next(self):
    return self._input_data

239
240
241
242
243
244
245
246
247
  def next(self):
    return self.__next__()

  def __next__(self):
    try:
      return self.get_next()
    except tf.errors.OutOfRangeError:
      raise StopIteration

248
249
250
251
252
253
254
255
256
  def initialize(self):
    if tf.executing_eagerly():
      return tf.no_op()
    else:
      return self._initializers


def _monkey_patch_dataset_method(strategy):
  """Monkey-patch `strategy`'s `make_dataset_iterator` method."""
257
  def make_dataset(self, dataset):
Toby Boyd's avatar
Toby Boyd committed
258
    tf.compat.v1.logging.info('Using pure synthetic data.')
259
260
261
262
263
264
    with self.scope():
      if self.extended._global_batch_size:  # pylint: disable=protected-access
        return SyntheticDataset(dataset, self.num_replicas_in_sync)
      else:
        return SyntheticDataset(dataset)

265
266
267
268
269
270
271
272
  def make_iterator(self, dataset):
    dist_dataset = make_dataset(self, dataset)
    return iter(dist_dataset)

  strategy.orig_make_dataset_iterator = strategy.make_dataset_iterator
  strategy.make_dataset_iterator = make_iterator
  strategy.orig_distribute_dataset = strategy.experimental_distribute_dataset
  strategy.experimental_distribute_dataset = make_dataset
273
274
275


def _undo_monkey_patch_dataset_method(strategy):
276
277
278
279
  if hasattr(strategy, 'orig_make_dataset_iterator'):
    strategy.make_dataset_iterator = strategy.orig_make_dataset_iterator
  if hasattr(strategy, 'orig_distribute_dataset'):
    strategy.make_dataset_iterator = strategy.orig_distribute_dataset
280
281
282


def set_up_synthetic_data():
283
  _monkey_patch_dataset_method(tf.distribute.OneDeviceStrategy)
284
  _monkey_patch_dataset_method(tf.distribute.MirroredStrategy)
285
286
  _monkey_patch_dataset_method(
      tf.distribute.experimental.MultiWorkerMirroredStrategy)
Toby Boyd's avatar
Toby Boyd committed
287
288
  # TODO(tobyboyd): Remove when contrib.distribute is all in core.
  if hasattr(tf, 'contrib'):
289
290
291
    _monkey_patch_dataset_method(contrib_distribute.MirroredStrategy)
    _monkey_patch_dataset_method(contrib_distribute.OneDeviceStrategy)
    _monkey_patch_dataset_method(contrib_distribute.CollectiveAllReduceStrategy)
Toby Boyd's avatar
Toby Boyd committed
292
293
  else:
    print('Contrib missing: Skip monkey patch tf.contrib.distribute.*')
294
295
296


def undo_set_up_synthetic_data():
297
  _undo_monkey_patch_dataset_method(tf.distribute.OneDeviceStrategy)
298
  _undo_monkey_patch_dataset_method(tf.distribute.MirroredStrategy)
299
300
  _undo_monkey_patch_dataset_method(
      tf.distribute.experimental.MultiWorkerMirroredStrategy)
Toby Boyd's avatar
Toby Boyd committed
301
302
  # TODO(tobyboyd): Remove when contrib.distribute is all in core.
  if hasattr(tf, 'contrib'):
303
304
    _undo_monkey_patch_dataset_method(contrib_distribute.MirroredStrategy)
    _undo_monkey_patch_dataset_method(contrib_distribute.OneDeviceStrategy)
305
    _undo_monkey_patch_dataset_method(
306
        contrib_distribute.CollectiveAllReduceStrategy)
Toby Boyd's avatar
Toby Boyd committed
307
308
  else:
    print('Contrib missing: Skip remove monkey patch tf.contrib.distribute.*')
309
310
311
312
313
314
315
316
317
318
319
320
321


def configure_cluster(worker_hosts=None, task_index=-1):
  """Set multi-worker cluster spec in TF_CONFIG environment variable.

  Args:
    worker_hosts: comma-separated list of worker ip:port pairs.

  Returns:
    Number of workers in the cluster.
  """
  tf_config = json.loads(os.environ.get('TF_CONFIG', '{}'))
  if tf_config:
322
323
    num_workers = (len(tf_config['cluster'].get('chief', [])) +
                   len(tf_config['cluster'].get('worker', [])))
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
  elif worker_hosts:
    workers = worker_hosts.split(',')
    num_workers = len(workers)
    if num_workers > 1 and task_index < 0:
      raise ValueError('Must specify task_index when number of workers > 1')
    task_index = 0 if num_workers == 1 else task_index
    os.environ['TF_CONFIG'] = json.dumps({
        'cluster': {
            'worker': workers
        },
        'task': {'type': 'worker', 'index': task_index}
    })
  else:
    num_workers = 1
  return num_workers
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356


def get_strategy_scope(strategy):
  if strategy:
    strategy_scope = strategy.scope()
  else:
    strategy_scope = DummyContextManager()

  return strategy_scope


class DummyContextManager(object):

  def __enter__(self):
    pass

  def __exit__(self, *args):
    pass