Makes new resource handle ops corresponding to existing resource tensors. Creates resource handle ops in the current default graph, whereas `accessible_objects` will be from an eager context. Resource mapping adds resource handle ops to the main GraphDef of a SavedModel, which allows th
(self)
| 214 | child_proto.local_name = local_name |
| 215 | |
| 216 | def map_resources(self): |
| 217 | """Makes new resource handle ops corresponding to existing resource tensors. |
| 218 | |
| 219 | Creates resource handle ops in the current default graph, whereas |
| 220 | `accessible_objects` will be from an eager context. Resource mapping adds |
| 221 | resource handle ops to the main GraphDef of a SavedModel, which allows the |
| 222 | C++ loader API to interact with variables. |
| 223 | |
| 224 | Returns: |
| 225 | A tuple of (object_map, resource_map, asset_info): |
| 226 | object_map: A dictionary mapping from object in `accessible_objects` to |
| 227 | replacement objects created to hold the new resource tensors. |
| 228 | resource_map: A dictionary mapping from resource tensors extracted from |
| 229 | `accessible_objects` to newly created resource tensors. |
| 230 | asset_info: An _AssetInfo tuple describing external assets referenced |
| 231 | from accessible_objects. |
| 232 | """ |
| 233 | # Only makes sense when adding to the export Graph |
| 234 | assert not context.executing_eagerly() |
| 235 | # TODO(allenl): Handle MirroredVariables and other types of variables which |
| 236 | # may need special casing. |
| 237 | object_map = object_identity.ObjectIdentityDictionary() |
| 238 | resource_map = {} |
| 239 | asset_info = _AssetInfo( |
| 240 | asset_defs=[], |
| 241 | asset_initializers_by_resource={}, |
| 242 | asset_filename_map={}, |
| 243 | asset_index={}) |
| 244 | |
| 245 | for node_id, obj in enumerate(self.nodes): |
| 246 | if isinstance(obj, tracking.CapturableResource): |
| 247 | # pylint: disable=protected-access |
| 248 | with ops.device(obj._resource_device): |
| 249 | new_resource = obj._create_resource() |
| 250 | # pylint: enable=protected-access |
| 251 | resource_map[obj.resource_handle] = new_resource |
| 252 | self.captured_tensor_node_ids[obj.resource_handle] = node_id |
| 253 | elif ds_values.is_distributed_variable(obj): |
| 254 | # Put both the distributed variable and component variable handles in |
| 255 | # `captured_tensor_node_ids`. |
| 256 | # Also create a new distributed variable for `object_map` with newly |
| 257 | # created component variables. |
| 258 | new_vars = [] |
| 259 | for v in obj.values: |
| 260 | new_variable = resource_variable_ops.copy_to_graph_uninitialized(v) |
| 261 | object_map[v] = new_variable |
| 262 | new_vars.append(new_variable) |
| 263 | resource_map[v.handle] = new_variable.handle |
| 264 | self.captured_tensor_node_ids[v.handle] = node_id |
| 265 | object_map[obj] = obj._clone_with_new_values(new_vars) # pylint: disable=protected-access |
| 266 | self.captured_tensor_node_ids[obj] = node_id |
| 267 | elif resource_variable_ops.is_resource_variable(obj): |
| 268 | new_variable = resource_variable_ops.copy_to_graph_uninitialized(obj) |
| 269 | object_map[obj] = new_variable |
| 270 | resource_map[obj.handle] = new_variable.handle |
| 271 | self.captured_tensor_node_ids[obj.handle] = node_id |
| 272 | elif isinstance(obj, tracking.TrackableAsset): |
| 273 | _process_asset(obj, asset_info, resource_map) |
no test coverage detected