Load graphs and restore variable values from a `SavedModel`.
| 270 | |
| 271 | |
| 272 | class SavedModelLoader(object): |
| 273 | """Load graphs and restore variable values from a `SavedModel`.""" |
| 274 | |
| 275 | def __init__(self, export_dir): |
| 276 | """Creates a `SavedModelLoader`. |
| 277 | |
| 278 | Args: |
| 279 | export_dir: Directory in which the SavedModel protocol buffer and |
| 280 | variables to be loaded are located. |
| 281 | """ |
| 282 | self._export_dir = export_dir |
| 283 | self._variables_path = saved_model_utils.get_variables_path(export_dir) |
| 284 | self._saved_model = parse_saved_model(export_dir) |
| 285 | |
| 286 | @property |
| 287 | def export_dir(self): |
| 288 | """Directory containing the SavedModel.""" |
| 289 | return self._export_dir |
| 290 | |
| 291 | @property |
| 292 | def variables_path(self): |
| 293 | """Path to variable checkpoint files.""" |
| 294 | return self._variables_path |
| 295 | |
| 296 | @property |
| 297 | def saved_model(self): |
| 298 | """SavedModel object parsed from the export directory.""" |
| 299 | return self._saved_model |
| 300 | |
| 301 | def get_meta_graph_def_from_tags(self, tags): |
| 302 | """Return MetaGraphDef with the exact specified tags. |
| 303 | |
| 304 | Args: |
| 305 | tags: A list or set of string tags that identify the MetaGraphDef. |
| 306 | |
| 307 | Returns: |
| 308 | MetaGraphDef with the same tags. |
| 309 | |
| 310 | Raises: |
| 311 | RuntimeError: if no metagraphs were found with the associated tags. |
| 312 | """ |
| 313 | found_match = False |
| 314 | available_tags = [] |
| 315 | for meta_graph_def in self._saved_model.meta_graphs: |
| 316 | available_tags.append(set(meta_graph_def.meta_info_def.tags)) |
| 317 | if set(meta_graph_def.meta_info_def.tags) == set(tags): |
| 318 | meta_graph_def_to_load = meta_graph_def |
| 319 | found_match = True |
| 320 | break |
| 321 | |
| 322 | if not found_match: |
| 323 | raise RuntimeError( |
| 324 | "MetaGraphDef associated with tags " + str(tags).strip("[]") + |
| 325 | " could not be found in SavedModel. To inspect available tag-sets in" |
| 326 | " the SavedModel, please use the SavedModel CLI: `saved_model_cli`" |
| 327 | "\navailable_tags: " + str(available_tags)) |
| 328 | return meta_graph_def_to_load |
| 329 |