Returns an empty dataset indicator and the next input from the iterator.
(iterator, strategy, name=None)
| 184 | |
| 185 | |
| 186 | def _get_next_as_optional(iterator, strategy, name=None): |
| 187 | """Returns an empty dataset indicator and the next input from the iterator.""" |
| 188 | replicas = [] |
| 189 | worker_has_values = [] |
| 190 | worker_devices = [] |
| 191 | for i, worker in enumerate(iterator._input_workers.worker_devices): # pylint: disable=protected-access |
| 192 | if name is not None: |
| 193 | d = tf_device.DeviceSpec.from_string(worker) |
| 194 | new_name = "%s_%s_%d" % (name, d.job, d.task) |
| 195 | else: |
| 196 | new_name = None |
| 197 | |
| 198 | with ops.device(worker): |
| 199 | worker_has_value, next_element = ( |
| 200 | iterator._iterators[i].get_next_as_list(new_name)) # pylint: disable=protected-access |
| 201 | # Collective all-reduce requires explict devices for inputs. |
| 202 | with ops.device("/cpu:0"): |
| 203 | # Converting to integers for all-reduce. |
| 204 | worker_has_value = math_ops.cast(worker_has_value, dtypes.int32) |
| 205 | worker_devices.append(worker_has_value.device) |
| 206 | worker_has_values.append(worker_has_value) |
| 207 | # Make `replicas` a flat list of values across all replicas. |
| 208 | replicas.append(next_element) |
| 209 | |
| 210 | # Run an all-reduce to see whether any worker has values. |
| 211 | # TODO(b/131423105): we should be able to short-cut the all-reduce in some |
| 212 | # cases. |
| 213 | if getattr(strategy.extended, "_support_per_replica_values", True): |
| 214 | worker_has_values = values.PerReplica( |
| 215 | values.WorkerDeviceMap( |
| 216 | worker_devices, |
| 217 | num_replicas_per_worker=len( |
| 218 | strategy.extended._input_workers._input_worker_devices)), # pylint: disable=protected-access |
| 219 | worker_has_values) |
| 220 | global_has_value = strategy.reduce( |
| 221 | reduce_util.ReduceOp.SUM, worker_has_values, axis=None) |
| 222 | else: |
| 223 | assert len(worker_has_values) == 1 |
| 224 | global_has_value = worker_has_values[0] |
| 225 | global_has_value = array_ops.reshape( |
| 226 | math_ops.cast(global_has_value, dtypes.bool), []) |
| 227 | return global_has_value, replicas |
| 228 | |
| 229 | |
| 230 | class DistributedIterator(object): |
no test coverage detected