Builds saver_def.
(self, checkpoint_path, build_save, build_restore)
| 1056 | checkpoint_path, build_save=build_save, build_restore=build_restore) |
| 1057 | |
| 1058 | def _build(self, checkpoint_path, build_save, build_restore): |
| 1059 | """Builds saver_def.""" |
| 1060 | if not context.executing_eagerly(): |
| 1061 | if self._is_built: |
| 1062 | return |
| 1063 | self._is_built = True |
| 1064 | |
| 1065 | if not self.saver_def or context.executing_eagerly(): |
| 1066 | if self._builder is None: |
| 1067 | self._builder = BulkSaverBuilder(self._write_version, self._incremental_save_restore, |
| 1068 | self._incremental_include_normal_var) |
| 1069 | |
| 1070 | if self._var_list is None: |
| 1071 | # pylint: disable=protected-access |
| 1072 | self._var_list = variables._all_saveable_objects() |
| 1073 | from tensorflow.python.ops import hash_table |
| 1074 | from tensorflow.python.ops import kv_variable_ops |
| 1075 | if isinstance(self._var_list, dict): |
| 1076 | ev = {} |
| 1077 | ht = {} |
| 1078 | lst = {} |
| 1079 | for name, x in self._var_list.items(): |
| 1080 | if isinstance(x, kv_variable_ops.EmbeddingVariable): |
| 1081 | ev[name] = x |
| 1082 | if isinstance(x, hash_table.HashTable): |
| 1083 | if x.hash_table not in ht: |
| 1084 | ht[x.hash_table] = [x] |
| 1085 | else: |
| 1086 | ht[x.hash_table].append(x) |
| 1087 | elif isinstance(x, hash_table.BloomFilterAdmitStrategy): |
| 1088 | lst[name] = BloomFilterSaveable(x) |
| 1089 | else: |
| 1090 | lst[name] = x |
| 1091 | if len(ev) != 0 and not self._sharded: |
| 1092 | raise ValueError("EmbeddingVariable can only use sharded saver") |
| 1093 | if len(ht) != 0 and not self._sharded: |
| 1094 | raise ValueError("HashTable can only use sharded saver") |
| 1095 | for x, y in ht.items(): |
| 1096 | lst[x.name] = HashTableSaveable(y) |
| 1097 | self._var_list = lst |
| 1098 | else: |
| 1099 | ev = [] |
| 1100 | ht = {} |
| 1101 | lst = [] |
| 1102 | for x in self._var_list: |
| 1103 | if isinstance(x, kv_variable_ops.EmbeddingVariable): |
| 1104 | ev.append(x) |
| 1105 | if isinstance(x, hash_table.HashTable): |
| 1106 | if x.hash_table not in ht: |
| 1107 | ht[x.hash_table] = [x] |
| 1108 | else: |
| 1109 | ht[x.hash_table].append(x) |
| 1110 | elif isinstance(x, hash_table.BloomFilterAdmitStrategy): |
| 1111 | lst.append(BloomFilterSaveable(x)) |
| 1112 | else: |
| 1113 | lst.append(x) |
| 1114 | if len(ev) != 0 and not self._sharded: |
| 1115 | raise ValueError("EmbeddingVariable can only use sharded saver") |
no test coverage detected