Adds tensors from potentially multiple devices.
(tensor_list, gradient_uid)
| 880 | |
| 881 | |
| 882 | def _MultiDeviceAddN(tensor_list, gradient_uid): |
| 883 | """Adds tensors from potentially multiple devices.""" |
| 884 | # Basic function structure comes from control_flow_ops.group(). |
| 885 | # Sort tensors according to their devices. |
| 886 | tensors_on_device = collections.defaultdict(lambda: []) |
| 887 | for tensor in tensor_list: |
| 888 | tensors_on_device[tensor.device].append(tensor) |
| 889 | |
| 890 | # For each device, add the tensors on that device first. |
| 891 | # Then gather the partial sums from multiple devices. |
| 892 | # TODO(sjhwang): Create hierarchical aggregation tree as pbar's suggestion. |
| 893 | # E.g., aggregate per GPU, then per task, and so on. |
| 894 | summands = [] |
| 895 | |
| 896 | def DeviceKey(dev): |
| 897 | return "" if dev is None else dev |
| 898 | |
| 899 | for dev in sorted(tensors_on_device, key=DeviceKey): |
| 900 | tensors = tensors_on_device[dev] |
| 901 | with ops._colocate_with_for_gradient( # pylint: disable=protected-access |
| 902 | tensors[0].op, |
| 903 | gradient_uid, |
| 904 | ignore_existing=True): |
| 905 | summands.append(math_ops.add_n(tensors)) |
| 906 | |
| 907 | return math_ops.add_n(summands) |
| 908 | |
| 909 | |
| 910 | @tf_export("AggregationMethod") |
no test coverage detected