Finds resource variables and lifts them into the outer context. When we import a GraphDef inside a wrap_function, no Python graph building code runs. This means we get VarHandleOps which create variable resources, but no corresponding Python objects. Leaving them like this works but gives t
(graph, variable_holder)
| 143 | |
| 144 | |
| 145 | def _lift_unlifted_variables(graph, variable_holder): |
| 146 | """Finds resource variables and lifts them into the outer context. |
| 147 | |
| 148 | When we import a GraphDef inside a wrap_function, no Python graph building |
| 149 | code runs. This means we get VarHandleOps which create variable resources, |
| 150 | but no corresponding Python objects. Leaving them like this works but gives |
| 151 | the user no way to interact with or modify the variables outside the graph. |
| 152 | |
| 153 | This method searches for variables and lifts them out as regular variable |
| 154 | objects when possible, indicating to the FuncGraph that they are captures. |
| 155 | |
| 156 | Args: |
| 157 | graph: The FuncGraph to lift variables from. |
| 158 | variable_holder: A VariableHolder to record the lifted variables in. |
| 159 | """ |
| 160 | with graph.as_default(): |
| 161 | global_collection_variables = ops.get_collection( |
| 162 | ops.GraphKeys.GLOBAL_VARIABLES) |
| 163 | local_collection_variables = ops.get_collection( |
| 164 | ops.GraphKeys.LOCAL_VARIABLES) |
| 165 | existing_captures = object_identity.ObjectIdentitySet( |
| 166 | graph.internal_captures) |
| 167 | lifted_variables = object_identity.ObjectIdentityDictionary() |
| 168 | |
| 169 | def _should_lift_variable(v): |
| 170 | return ((v._in_graph_mode # pylint: disable=protected-access |
| 171 | and v.graph.building_function) |
| 172 | and isinstance(v, resource_variable_ops.BaseResourceVariable) |
| 173 | and v.handle not in existing_captures) |
| 174 | |
| 175 | for old_variable in global_collection_variables: |
| 176 | if _should_lift_variable(old_variable): |
| 177 | new_variable = _lift_single_variable( |
| 178 | old_variable, graph, variable_holder) |
| 179 | lifted_variables[old_variable] = new_variable |
| 180 | existing_captures.add(old_variable.handle) |
| 181 | |
| 182 | for old_variable in local_collection_variables: |
| 183 | if _should_lift_variable(old_variable): |
| 184 | new_variable = _lift_single_variable( |
| 185 | old_variable, graph, variable_holder) |
| 186 | lifted_variables[old_variable] = new_variable |
| 187 | existing_captures.add(old_variable.handle) |
| 188 | if new_variable._in_graph_mode: # pylint: disable=protected-access |
| 189 | outer_graph = new_variable.graph |
| 190 | # Variables are added to the global collection by default. In this |
| 191 | # case we only want the variable in the local collection, so we'll pop |
| 192 | # it out. |
| 193 | global_collection = outer_graph.get_collection_ref( |
| 194 | ops.GraphKeys.GLOBAL_VARIABLES) |
| 195 | global_collection.remove(new_variable) |
| 196 | outer_graph.add_to_collection( |
| 197 | ops.GraphKeys.LOCAL_VARIABLES, new_variable) |
| 198 | |
| 199 | # Update the FuncGraph's collections, partly for the user and partly so this |
| 200 | # function is idempotent when it runs again in prune() calls. |
| 201 | for collection_name in [ |
| 202 | ops.GraphKeys.GLOBAL_VARIABLES, ops.GraphKeys.LOCAL_VARIABLES |
no test coverage detected