AllReduce the gradients with raw ops (instead of collective ops). Args: all_grads (K x N): A list of K lists. Each of the list is a list of N grad tuples. The variables have to be the same across the K lists. devices (list[str]): assign the averaging to these de
(all_grads, devices=None, average=True)
| 295 | |
| 296 | @under_name_scope('AllReduceNaive') |
| 297 | def allreduce_grads_naive(all_grads, devices=None, average=True): |
| 298 | """ |
| 299 | AllReduce the gradients with raw ops (instead of collective ops). |
| 300 | |
| 301 | Args: |
| 302 | all_grads (K x N): A list of K lists. Each of the list is a list of N grad tuples. |
| 303 | The variables have to be the same across the K lists. |
| 304 | devices (list[str]): assign the averaging to these device in |
| 305 | round-robin. Cannot be used together with ``colocation``. |
| 306 | average (bool): do average or sum |
| 307 | |
| 308 | Returns: |
| 309 | list[Tensor]: list of grads where each grad is averaged or summed over K. |
| 310 | """ |
| 311 | if devices is not None: |
| 312 | assert isinstance(devices, list), devices |
| 313 | # device_setter = LeastLoadedDeviceSetter(None, devices) |
| 314 | |
| 315 | nr_tower = len(all_grads) |
| 316 | if nr_tower == 1: |
| 317 | return all_grads[0] |
| 318 | |
| 319 | def aggregate(grads): |
| 320 | if average: |
| 321 | return tf.multiply(tf.add_n(grads), 1.0 / nr_tower) |
| 322 | else: |
| 323 | return tf.add_n(grads) |
| 324 | |
| 325 | grads_ret = [] # N(rev) grads |
| 326 | # reverse so the device placement makes the last part of model more balance? |
| 327 | all_grads_rev = [x[::-1] for x in all_grads] # K x N(rev) |
| 328 | |
| 329 | for idx, grads in enumerate(zip(*all_grads_rev)): |
| 330 | # grads: K tensors |
| 331 | if devices is None: |
| 332 | grad = aggregate(grads) |
| 333 | else: |
| 334 | # dev = device_setter.place_with_balance(v.op) |
| 335 | dev = devices[idx % len(devices)] |
| 336 | with tf.device(dev): |
| 337 | grad = aggregate(grads) |
| 338 | grads_ret.append(grad) |
| 339 | grads_ret = grads_ret[::-1] |
| 340 | return grads_ret |
| 341 | |
| 342 | |
| 343 | # https://github.com/tensorflow/benchmarks/blob/48cbef14a592e02a14beee8e9aef3ad22cadaed1/scripts/tf_cnn_benchmarks/variable_mgr_util.py#L140-L166 |