(self, object_graph_proto, saved_model_proto, export_dir)
| 103 | """Helper class to load an object-based SavedModel.""" |
| 104 | |
| 105 | def __init__(self, object_graph_proto, saved_model_proto, export_dir): |
| 106 | meta_graph = saved_model_proto.meta_graphs[0] |
| 107 | self._asset_file_def = meta_graph.asset_file_def |
| 108 | self._operation_attributes = { |
| 109 | node.name: node.attr for node in meta_graph.graph_def.node} |
| 110 | self._proto = object_graph_proto |
| 111 | self._export_dir = export_dir |
| 112 | self._concrete_functions = ( |
| 113 | function_deserialization.load_function_def_library( |
| 114 | meta_graph.graph_def.library)) |
| 115 | |
| 116 | for name, concrete_function in self._concrete_functions.items(): |
| 117 | # Wrap all the concrete function so that they are capable of dealing with |
| 118 | # both in replica and cross replica cases. |
| 119 | self._concrete_functions[name] = _WrapperFunction(concrete_function) |
| 120 | |
| 121 | self._load_all() |
| 122 | # TODO(b/124045874): There are limitations with functions whose captures |
| 123 | # trigger other functions to be executed. For now it is only guaranteed to |
| 124 | # work if the captures of a function only trigger functions without |
| 125 | # captures. |
| 126 | self._setup_functions_structures() |
| 127 | self._setup_functions_captures() |
| 128 | self._restore_checkpoint() |
| 129 | |
| 130 | for node in self._nodes: |
| 131 | if isinstance(node, tracking.CapturableResource): |
| 132 | init_op = node._initialize() # pylint: disable=protected-access |
| 133 | if not context.executing_eagerly(): |
| 134 | ops.add_to_collection(ops.GraphKeys.TABLE_INITIALIZERS, init_op) |
| 135 | |
| 136 | def _setup_functions_structures(self): |
| 137 | """Setup structure for inputs and outputs of restored functions.""" |
no test coverage detected