SaverBuilder with support for bulk restoring multiple saveables.
| 715 | |
| 716 | |
| 717 | class BulkSaverBuilder(BaseSaverBuilder): |
| 718 | """SaverBuilder with support for bulk restoring multiple saveables.""" |
| 719 | |
| 720 | def bulk_restore(self, filename_tensor, saveables, preferred_shard, |
| 721 | restore_sequentially): |
| 722 | |
| 723 | # Ignored: bulk restore is internally sequential. |
| 724 | del restore_sequentially |
| 725 | restore_specs = [] |
| 726 | for saveable in saveables: |
| 727 | if isinstance(saveable, BaseSaverBuilder.EmbeddingVariableSaveable): |
| 728 | continue |
| 729 | if saveable.custom_restore: |
| 730 | continue |
| 731 | for spec in saveable.specs: |
| 732 | restore_specs.append((spec.name, spec.slice_spec, spec.dtype)) |
| 733 | |
| 734 | if len(restore_specs) > 0: |
| 735 | names, slices, dtypes = zip(*restore_specs) |
| 736 | # Load all tensors onto CPU 0 for compatibility with existing code. |
| 737 | with ops.device("cpu:0"): |
| 738 | return io_ops.restore_v2(filename_tensor, names, slices, dtypes) |
| 739 | |
| 740 | class PartialRestoreSaverBuilder(BulkSaverBuilder): |
| 741 | """SaverBuilder with support for custom save and restore. |
no outgoing calls
no test coverage detected