Group Variable tensor slices per device. TODO(touts): Make sure that all the devices found are on different job/replica/task/cpu|gpu. It would be bad if 2 were on the same device. It can happen if the devices are unspecified. Args: saveables: A list of BaseSaverBuilder.Savea
(self, saveables)
| 532 | return control_flow_ops.group(*sharded_restores, name="restore_all") |
| 533 | |
| 534 | def _GroupByDevices(self, saveables): |
| 535 | """Group Variable tensor slices per device. |
| 536 | |
| 537 | TODO(touts): Make sure that all the devices found are on different |
| 538 | job/replica/task/cpu|gpu. It would be bad if 2 were on the same device. |
| 539 | It can happen if the devices are unspecified. |
| 540 | |
| 541 | Args: |
| 542 | saveables: A list of BaseSaverBuilder.SaveableObject objects. |
| 543 | |
| 544 | Returns: |
| 545 | A list of tuples: (device_name, BaseSaverBuilder.SaveableObject) tuples. |
| 546 | The list is sorted by ascending device_name. |
| 547 | |
| 548 | Raises: |
| 549 | ValueError: If the tensors of a saveable are on different devices. |
| 550 | """ |
| 551 | per_device = collections.defaultdict(lambda: []) |
| 552 | for saveable in saveables: |
| 553 | canonical_device = set() |
| 554 | for spec in saveable.specs: |
| 555 | device_name = pydev.canonical_name(spec.tensor.device) |
| 556 | device_spec = pydev.DeviceSpec.from_string(device_name) |
| 557 | device_spec.device_type = "CPU" |
| 558 | canonical_device.add(device_spec.to_string()) |
| 559 | if len(canonical_device) != 1: |
| 560 | raise ValueError("All tensors of a saveable object must be " |
| 561 | "on the same device: %s" % saveable.name) |
| 562 | per_device[canonical_device.pop()].append(saveable) |
| 563 | return sorted(per_device.items(), key=lambda t: t[0]) |
| 564 | |
| 565 | def build(self, |
| 566 | names_to_saveables, |
no test coverage detected