(self,
container_strategy,
tpu_cluster_resolver=None,
steps_per_run=None,
device_assignment=None)
| 155 | """Implementation of TPUStrategy.""" |
| 156 | |
| 157 | def __init__(self, |
| 158 | container_strategy, |
| 159 | tpu_cluster_resolver=None, |
| 160 | steps_per_run=None, |
| 161 | device_assignment=None): |
| 162 | super(TPUExtended, self).__init__(container_strategy) |
| 163 | |
| 164 | if tpu_cluster_resolver is None: |
| 165 | tpu_cluster_resolver = TPUClusterResolver("") |
| 166 | |
| 167 | if steps_per_run is None: |
| 168 | # TODO(frankchn): Warn when we are being used by DS/Keras and this is |
| 169 | # not specified. |
| 170 | steps_per_run = 1 |
| 171 | |
| 172 | self._tpu_function_cache = weakref.WeakKeyDictionary() |
| 173 | self._tpu_cluster_resolver = tpu_cluster_resolver |
| 174 | self._tpu_metadata = get_tpu_system_metadata(self._tpu_cluster_resolver) |
| 175 | self._device_assignment = device_assignment |
| 176 | |
| 177 | # Device assignment is currently only supported for 1 core case. |
| 178 | if self._device_assignment: |
| 179 | assert isinstance(self._device_assignment, |
| 180 | device_assignment_lib.DeviceAssignment) |
| 181 | if self._device_assignment.num_replicas != 1: |
| 182 | raise ValueError("Device assignment is only supported for a single " |
| 183 | "core single replica case currently.") |
| 184 | if self._device_assignment.num_cores_per_replica != 1: |
| 185 | raise ValueError("Device assignment is only supported for a single " |
| 186 | "core single replica case currently.") |
| 187 | if not all(self._device_assignment.core_assignment[0][0] == [0, 0, 0]): |
| 188 | raise ValueError("Device assignment is only supported for a single " |
| 189 | "core single replica case currently.") |
| 190 | |
| 191 | # TODO(jhseu): Switch to DeviceAssignment to support pods and model |
| 192 | # parallelism. |
| 193 | self._tpu_devices = [d.name for d in self._tpu_metadata.devices |
| 194 | if "device:TPU:" in d.name] |
| 195 | |
| 196 | self._host_device = device_util.get_host_for_device(self._tpu_devices[0]) |
| 197 | |
| 198 | # Only create variables for the number of replicas we're running. |
| 199 | self._tpu_devices = self._tpu_devices[:self._num_replicas_in_sync] |
| 200 | self._device_map = values.ReplicaDeviceMap(self._tpu_devices) |
| 201 | |
| 202 | # Preload the data onto the TPUs. |
| 203 | input_worker_devices = collections.OrderedDict() |
| 204 | for tpu_device in self._tpu_devices: |
| 205 | host_device = device_util.get_host_for_device(tpu_device) |
| 206 | input_worker_devices.setdefault(host_device, []) |
| 207 | input_worker_devices[host_device].append(tpu_device) |
| 208 | self._input_workers = input_lib.InputWorkers( |
| 209 | self._device_map, tuple(input_worker_devices.items())) |
| 210 | |
| 211 | # TODO(sourabhbajaj): Remove this once performance of running one step |
| 212 | # at a time is comparable to multiple steps. |
| 213 | self.steps_per_run = steps_per_run |
| 214 | self._require_static_shapes = True |
nothing calls this directly
no test coverage detected