Args: gpus (int or [int]): list of GPU ids. average (bool): whether to average or sum gradients. mode (str or None): Gradient aggregation mode. Supported values: ['nccl', 'hierarchical', 'cpu', 'gpu']. These modes may diffe
(self, gpus, average=True, mode=None)
| 167 | |
| 168 | @map_arg(gpus=_int_to_range) |
| 169 | def __init__(self, gpus, average=True, mode=None): |
| 170 | """ |
| 171 | Args: |
| 172 | gpus (int or [int]): list of GPU ids. |
| 173 | average (bool): whether to average or sum gradients. |
| 174 | mode (str or None): Gradient aggregation mode. |
| 175 | Supported values: ['nccl', 'hierarchical', 'cpu', 'gpu']. |
| 176 | These modes may differ in speed. |
| 177 | Default to pick automatically by heuristics. |
| 178 | "hierarchical" mode was designed for DGX-like 8-GPU machines. |
| 179 | """ |
| 180 | self.devices = gpus |
| 181 | if mode is not None: |
| 182 | mode = mode.lower() |
| 183 | |
| 184 | # Heuristics about mode selection: |
| 185 | if mode == 'hierarchical' and len(gpus) != 8: |
| 186 | logger.warn("mode='hierarchical' requires 8 GPUs. Will fallback to default mode.") |
| 187 | mode = None |
| 188 | if mode is None: |
| 189 | if len(gpus) == 8: |
| 190 | mode = 'hierarchical' |
| 191 | else: |
| 192 | # https://github.com/tensorflow/tensorflow/issues/41539 |
| 193 | mode = 'nccl' if get_tf_version_tuple() < (1, 15) else 'gpu' |
| 194 | if mode == 'cpu' and get_tf_version_tuple() >= (2, 0): |
| 195 | # cpu mode causes the entire model to get located on cpu |
| 196 | mode = 'gpu' |
| 197 | if mode == 'nccl' and get_tf_version_tuple() >= (1, 15): |
| 198 | logger.warning( |
| 199 | "NCCL in TensorFlow has a serious bug that is likely to trigger in TF>=1.15. " |
| 200 | "Try 'mode=None' to use a better default mode.") |
| 201 | |
| 202 | self._builder = SyncMultiGPUReplicatedBuilder(gpus, average, mode) |
| 203 | self.BROADCAST_EVERY_EPOCH = True |
| 204 | |
| 205 | super(SyncMultiGPUTrainerReplicated, self).__init__() |
| 206 | |
| 207 | def _setup_graph(self, input, get_cost_fn, get_opt_fn): |
| 208 | if len(self.devices) > 1: |
nothing calls this directly
no test coverage detected