Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py: 19%

1# Copyright 2018 The TensorFlow Authors. All Rights Reserved. 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================== 

15"""Exports a SavedModel from a Trackable Python object.""" 

16 

17import collections 

18import os 

19import re 

20import sys 

21import traceback 

22 

23from absl import logging 

24 

25from tensorflow.core.framework import function_pb2 

26from tensorflow.core.framework import versions_pb2 

27from tensorflow.core.protobuf import meta_graph_pb2 

28from tensorflow.core.protobuf import saved_model_pb2 

29from tensorflow.core.protobuf import saved_object_graph_pb2 

30from tensorflow.python.checkpoint import checkpoint 

31from tensorflow.python.checkpoint import checkpoint_options 

32from tensorflow.python.checkpoint import functional_saver 

33from tensorflow.python.checkpoint import graph_view 

34from tensorflow.python.checkpoint import save_util_v1 

35from tensorflow.python.checkpoint import util as checkpoint_util 

36from tensorflow.python.eager import context 

37from tensorflow.python.eager import def_function 

38from tensorflow.python.eager import function as defun 

39from tensorflow.python.eager.polymorphic_function import polymorphic_function 

40from tensorflow.python.eager.polymorphic_function import saved_model_exported_concrete 

41from tensorflow.python.eager.polymorphic_function import saved_model_utils 

42from tensorflow.python.framework import dtypes 

43from tensorflow.python.framework import error_interpolation 

44from tensorflow.python.framework import errors 

45from tensorflow.python.framework import function as framework_fn 

46from tensorflow.python.framework import meta_graph 

47from tensorflow.python.framework import ops 

48from tensorflow.python.framework import tensor_util 

49from tensorflow.python.framework import versions 

50from tensorflow.python.lib.io import file_io 

51from tensorflow.python.ops import array_ops 

52from tensorflow.python.ops import control_flow_ops 

53from tensorflow.python.ops import resource_variable_ops 

54from tensorflow.python.saved_model import builder_impl 

55from tensorflow.python.saved_model import fingerprinting_utils 

56from tensorflow.python.saved_model import function_serialization 

57from tensorflow.python.saved_model import path_helpers 

58from tensorflow.python.saved_model import pywrap_saved_model 

59from tensorflow.python.saved_model import registration 

60from tensorflow.python.saved_model import revived_types 

61from tensorflow.python.saved_model import save_context 

62from tensorflow.python.saved_model import save_options 

63from tensorflow.python.saved_model import signature_constants 

64from tensorflow.python.saved_model import signature_def_utils 

65from tensorflow.python.saved_model import signature_serialization 

66from tensorflow.python.saved_model import tag_constants 

67from tensorflow.python.saved_model import tracing_utils 

68from tensorflow.python.saved_model import utils_impl 

69from tensorflow.python.saved_model.pywrap_saved_model import constants 

70from tensorflow.python.saved_model.pywrap_saved_model import metrics 

71from tensorflow.python.trackable import asset 

72from tensorflow.python.trackable import base 

73from tensorflow.python.trackable import resource 

74from tensorflow.python.trackable import trackable_utils 

75from tensorflow.python.training.saving import trace_saveable_util 

76from tensorflow.python.types import core as types_core 

77from tensorflow.python.util import compat 

78from tensorflow.python.util import object_identity 

79from tensorflow.python.util.tf_export import tf_export 

80 

81_UNCOPIABLE_DTYPES = frozenset((dtypes.resource, dtypes.variant)) 

82 

83# Container for tensors captured from external functions. 

84_CapturedTensor = collections.namedtuple("_CapturedTensor", 

85 ["name", "concrete_function"]) 

86 

87# Number of untraced functions to display to user in warning message. 

88_NUM_DISPLAY_UNTRACED_FUNCTIONS = 5 

89 

90# API label for SavedModel metrics. 

91_SAVE_V2_LABEL = "save_v2" 

92 

93 

94class _AugmentedGraphView(graph_view.ObjectGraphView): 

95 """An extendable graph which also tracks functions attached to objects. 

96 

97 Extensions through `add_object` appear in the object graph and any checkpoints 

98 generated from it, even if they are not dependencies of the node they were 

99 attached to in the saving program. For example a `.signatures` attribute is 

100 added to exported SavedModel root objects without modifying the root object 

101 itself. 

102 

103 Also tracks functions attached to objects in the graph, through the caching 

104 `_list_functions` method. Enumerating functions only through this method 

105 ensures that we get a consistent view of functions, even if object attributes 

106 create new functions every time they are accessed. 

107 """ 

108 

109 def __init__(self, root): 

110 super(_AugmentedGraphView, self).__init__(root) 

111 

112 # Cache the results of `GraphView.list_children()` to ensure that the 

113 # `Trackable` children are gathered exactly once. 

114 self._children_cache = object_identity.ObjectIdentityDictionary() 

115 

116 # Cache shared between objects in the same object graph. This is passed to 

117 # `Trackable._trackable_children()`. 

118 self._serialization_cache = object_identity.ObjectIdentityDictionary() 

119 

120 # Maps functions -> wrapped functions that capture non-cached variables. 

121 self._wrapped_functions = {} 

122 

123 self.untraced_functions = [] 

124 

125 def set_signature(self, signature_map, wrapped_functions): 

126 """Attach signature to the root object. 

127 

128 Args: 

129 signature_map: An object that contains signature functions. 

130 wrapped_functions: A dictionary mapping functions to functions that are 

131 guaranteed to not capture cached variables (functions that capture 

132 cached variables can't be saved). 

133 """ 

134 self.list_children(self.root) 

135 # Overrides existing dependency. 

136 name = signature_serialization.SIGNATURE_ATTRIBUTE_NAME 

137 self._children_cache[self.root][name] = signature_map 

138 self._wrapped_functions.update(wrapped_functions) 

139 

140 def _breadth_first_traversal(self): 

141 """Returns all trackable objects in the SavedObjectGraph.""" 

142 # This method is overriden to merge all equivalent constant tensors and 

143 # Assets in the object graph. 

144 

145 trackable_objects, _ = ( 

146 super(_AugmentedGraphView, self)._breadth_first_traversal()) 

147 

148 asset_paths = object_identity.ObjectIdentityDictionary() 

149 constant_captures = object_identity.ObjectIdentityDictionary() 

150 for obj in trackable_objects: 

151 if isinstance(obj, asset.Asset): 

152 asset_paths[obj.asset_path] = obj 

153 if isinstance(obj, saved_model_utils.TrackableConstant): 

154 constant_captures[obj.capture] = obj 

155 

156 def _get_merged_trackable(x): 

157 if isinstance(x, asset.Asset): 

158 return asset_paths[x.asset_path] 

159 if isinstance(x, saved_model_utils.TrackableConstant): 

160 if x.capture in asset_paths: 

161 return asset_paths[x.capture] 

162 else: 

163 return constant_captures[x.capture] 

164 return x 

165 

166 for obj in list(self._children_cache.keys()): 

167 if _get_merged_trackable(obj) is not obj: 

168 del self._children_cache[obj] 

169 continue 

170 for name, child in self._children_cache[obj].items(): 

171 self._children_cache[obj][name] = _get_merged_trackable(child) 

172 

173 return super(_AugmentedGraphView, self)._breadth_first_traversal() 

174 

175 def list_children(self, obj): 

176 """Lists children of `obj` for SavedModel.""" 

177 if obj not in self._children_cache: 

178 children = self._children_cache[obj] = {} 

179 

180 for name, child in super(_AugmentedGraphView, self).list_children( 

181 obj, 

182 save_type=base.SaveType.SAVEDMODEL, 

183 cache=self._serialization_cache): 

184 if isinstance(child, defun.ConcreteFunction): 

185 child = self._maybe_uncache_variable_captures(child) 

186 children[name] = child 

187 

188 # Keep track of untraced functions for later reporting to the user. 

189 if isinstance(obj, def_function.Function) and not children: 

190 self.untraced_functions.append(obj.name) 

191 

192 for name, child in self._children_cache[obj].items(): 

193 yield base.TrackableReference(name, child) 

194 

195 def get_child(self, obj, name): 

196 return self._children_cache[obj][name] 

197 

198 def _maybe_uncache_variable_captures(self, concrete_function): 

199 if concrete_function in self._wrapped_functions: 

200 return self._wrapped_functions[concrete_function] 

201 for capture in concrete_function.captured_inputs: 

202 if hasattr(capture, "_cached_variable"): 

203 if concrete_function not in self._wrapped_functions: 

204 wrapped = self._wrapped_functions[concrete_function] = ( 

205 function_serialization.wrap_cached_variables(concrete_function)) 

206 return wrapped 

207 return concrete_function 

208 

209 def list_dependencies(self, obj): 

210 """Yields `Trackables` that must be loaded before `obj`. 

211 

212 Dependencies and children are both dictionaries of `Trackables`. Children 

213 define the object graph structure (used in both checkpoints and SavedModel), 

214 while dependency defines the order used to load the SavedModel 

215 

216 Args: 

217 obj: A `Trackable` object 

218 

219 Yields: 

220 Tuple of dependency names and trackable objects. 

221 

222 Raises: 

223 TypeError: if any of the returned dependencies are not instances of 

224 `Trackable`. 

225 """ 

226 if obj not in self._children_cache: 

227 # Slot variables do not appear in the children_cache. 

228 children = {} 

229 else: 

230 children = self._children_cache[obj] 

231 for name, dep in obj._deserialization_dependencies(children).items(): # pylint: disable=protected-access 

232 if not isinstance(dep, base.Trackable): 

233 raise TypeError( 

234 f"The dependency of type {type(dep)} is not an instance `Trackable`" 

235 ", and can't be saved to SavedModel. Please check the " 

236 "implementation of `_deserialization_dependencies` in the parent " 

237 f"object {obj}.") 

238 yield name, dep 

239 

240 

241class _SaveableView(object): 

242 """Provides a frozen view over a trackable root. 

243 

244 This class helps to create a single stable view over an object to save. The 

245 saving code should access properties and functions via this class and not via 

246 the original object as there are cases where an object construct their 

247 trackable attributes and functions dynamically per call and will yield 

248 different objects if invoked more than once. 

249 

250 Changes to the graph, for example adding objects, must happen in 

251 `augmented_graph_view` (an `_AugmentedGraphView`) before the `_SaveableView` 

252 is constructed. Changes after the `_SaveableView` has been constructed will be 

253 ignored. 

254 """ 

255 

256 def __init__(self, augmented_graph_view, options): 

257 """Initializes a SaveableView. 

258 

259 Args: 

260 augmented_graph_view: A GraphView object. 

261 options: A SaveOptions instance. 

262 """ 

263 

264 self.augmented_graph_view = augmented_graph_view 

265 self._options = options 

266 

267 (self._trackable_objects, self.node_paths, self.node_ids, 

268 self._slot_variables, self.object_names) = ( 

269 checkpoint_util.objects_ids_and_slot_variables_and_paths( 

270 self.augmented_graph_view)) 

271 

272 untraced_functions = self.augmented_graph_view.untraced_functions 

273 if untraced_functions: 

274 logging.info( 

275 "Found untraced functions such as %s while saving (showing %d of %d)." 

276 " These functions will not be directly callable after loading.", 

277 ", ".join(untraced_functions[:_NUM_DISPLAY_UNTRACED_FUNCTIONS]), 

278 min(_NUM_DISPLAY_UNTRACED_FUNCTIONS, len(untraced_functions)), 

279 len(untraced_functions)) 

280 

281 self._initialize_save_and_restore_functions() 

282 self._initialize_nodes_and_concrete_functions() 

283 

284 self.captured_tensor_node_ids = object_identity.ObjectIdentityDictionary() 

285 

286 def _initialize_save_and_restore_functions(self): 

287 """Generates all checkpoint save/restore functions. 

288 

289 The save and restore functions are generated in the eager context (or in the 

290 user's Graph/Session) before being copied to the exported GraphDef. These 

291 functions record the ops for saving/restoring the entire object or 

292 individual objects (e.g. variables and hash tables). 

293 

294 The global save and restore functions are generated for compatibility with 

295 TF1 and loading from C++, and is saved in the `MetaGraphDef.saver_def`. 

296 

297 The individual functions are generated for the Python TF2 use case, where 

298 users use the loaded SavedModel as-is, or compose new models using parts 

299 of the object loaded from the SavedModel. These functions are recorded in 

300 the `saveable_objects` map in the `SavedObject` proto. 

301 """ 

302 checkpoint_factory_map, registered_savers = ( 

303 save_util_v1.get_checkpoint_factories_and_keys(self.object_names)) 

304 self._obj_to_registered_saver = object_identity.ObjectIdentityDictionary() 

305 for saver_name, trackables in registered_savers.items(): 

306 for trackable in trackables.values(): 

307 self._obj_to_registered_saver[trackable] = saver_name 

308 self._saveable_objects_map = ( 

309 _gen_save_and_restore_functions(checkpoint_factory_map)) 

310 

311 def _initialize_nodes_and_concrete_functions(self): 

312 """Creates graph with nodes for trackable objects and functions. 

313 

314 Adds functions for each trackable object to `self.nodes` and associated 

315 concrete functions to `self.concrete_functions` for serialization. 

316 """ 

317 self.nodes = list(self._trackable_objects) 

318 self.gradient_functions = [] 

319 self.gradient_defs = [] 

320 

321 for obj in self.nodes: 

322 if obj in self._saveable_objects_map: 

323 for save_fn, restore_fn in self._saveable_objects_map[obj].values(): 

324 self.node_ids[save_fn] = len(self.nodes) 

325 self.nodes.append(save_fn) 

326 

327 self.node_ids[restore_fn] = len(self.nodes) 

328 self.nodes.append(restore_fn) 

329 

330 self.concrete_functions = [ 

331 obj for obj in self.nodes if isinstance(obj, defun.ConcreteFunction) 

332 ] 

333 

334 @property 

335 def concrete_and_gradient_functions(self): 

336 return self.concrete_functions + self.gradient_functions 

337 

338 @property 

339 def root(self): 

340 return self.nodes[0] 

341 

342 def fill_object_graph_proto(self, proto): 

343 """Populate the nodes, children and slot_variables of a SavedObjectGraph.""" 

344 for node_id, node in enumerate(self.nodes): 

345 assert self.node_ids[node] == node_id 

346 object_proto = proto.nodes.add() 

347 object_proto.slot_variables.extend(self._slot_variables.get(node, ())) 

348 if isinstance(node, _CapturedTensor): 

349 continue 

350 for child in self.augmented_graph_view.list_children(node): 

351 child_proto = object_proto.children.add() 

352 child_proto.node_id = self.node_ids[child.ref] 

353 child_proto.local_name = child.name 

354 for name, ref in self.augmented_graph_view.list_dependencies(node): 

355 child_proto = object_proto.dependencies.add() 

356 child_proto.node_id = self.node_ids[ref] 

357 child_proto.local_name = name 

358 

359 if node in self._saveable_objects_map: 

360 assert node not in self._obj_to_registered_saver, ( 

361 "Objects can't have both SaveableObjects and a registered saver") 

362 

363 for local_name, (save_fn, restore_fn) in ( 

364 self._saveable_objects_map[node].items()): 

365 saveable_object_proto = object_proto.saveable_objects[local_name] 

366 saveable_object_proto.save_function = self.node_ids[save_fn] 

367 saveable_object_proto.restore_function = self.node_ids[restore_fn] 

368 

369 elif node in self._obj_to_registered_saver: 

370 object_proto.registered_saver = self._obj_to_registered_saver[node] 

371 

372 def map_resources(self): 

373 """Makes new resource handle ops corresponding to existing resource tensors. 

374 

375 Creates resource handle ops in the current default graph, whereas 

376 `accessible_objects` will be from an eager context. Resource mapping adds 

377 resource handle ops to the main GraphDef of a SavedModel, which allows the 

378 C++ loader API to interact with resources. 

379 

380 Returns: 

381 A tuple of (object_map, tensor_map, asset_info): 

382 object_map: A dictionary mapping from object in `accessible_objects` to 

383 replacement objects created to hold the new resource tensors. 

384 tensor_map: A dictionary mapping from resource tensors extracted from 

385 `accessible_objects` to newly created resource tensors. 

386 asset_info: An _AssetInfo tuple describing external assets referenced 

387 from accessible_objects. 

388 """ 

389 # Only makes sense when adding to the export Graph 

390 assert not context.executing_eagerly() 

391 # TODO(b/205007558): Handle MirroredVariables and other types of variables 

392 # which may need special casing. 

393 object_map = object_identity.ObjectIdentityDictionary() 

394 tensor_map = object_identity.ObjectIdentityDictionary() 

395 asset_info = _AssetInfo( 

396 asset_defs=[], 

397 asset_initializers_by_resource=object_identity.ObjectIdentityDictionary(), 

398 asset_filename_map={}, 

399 asset_index={}) 

400 

401 for node_id in _dependency_sorted_node_ids(self): 

402 obj = self.nodes[node_id] 

403 tensors = obj._export_to_saved_model_graph( # pylint: disable=protected-access 

404 object_map=object_map, 

405 tensor_map=tensor_map, 

406 options=self._options) 

407 if isinstance(obj, asset.Asset): 

408 _add_asset_info(obj, asset_info, tensor_map[obj.asset_path]) 

409 if tensors: 

410 for tensor in tensors: 

411 self.captured_tensor_node_ids[tensor] = node_id 

412 

413 return object_map, tensor_map, asset_info 

414 

415 def add_capture_and_node(self, capture, node): 

416 node_id = len(self.nodes) 

417 self.nodes.append(node) 

418 self.node_ids[capture] = node_id 

419 self.node_ids[node] = node_id 

420 self.captured_tensor_node_ids[capture] = node_id 

421 return node_id 

422 

423 def get_concrete_resource_initializers(self): 

424 concrete_initializers = [] 

425 for obj in self.nodes: 

426 if isinstance(obj, resource.CapturableResource): 

427 concrete_initializers.append( 

428 self.augmented_graph_view.get_child( 

429 obj, "_initialize").get_concrete_function()) 

430 return concrete_initializers 

431 

432 

433def _gen_save_and_restore_functions(checkpoint_factory_map): 

434 """Generates global and individual save/restore concrete functions. 

435 

436 The global functions records the ops to save and restore the entire object to 

437 a file prefix, while the individual functions save and restore value tensors 

438 for resources. 

439 

440 This function is intended to run on the output of 

441 `save_util_v1.get_checkpoint_factories_and_keys(object_names)`, 

442 which returns the generated a map of `_CheckpointFactoryData`. 

443 

444 Args: 

445 checkpoint_factory_map: A dictionary mapping trackable objects to 

446 a list of `_CheckpointFactoryData`. 

447 

448 Returns: 

449 Tuple of ( 

450 saveable_fn_map: Maps obj -> factory name -> (concrete save, restore) 

451 ) 

452 """ 

453 # Maps obj -> factory attribute_name -> (concrete save, concrete restore) 

454 # This 

455 saveable_fn_map = object_identity.ObjectIdentityDictionary() 

456 

457 for obj, factory_data_list in checkpoint_factory_map.items(): 

458 if resource_variable_ops.is_resource_variable(obj) or not factory_data_list: 

459 # There is no need to trace the save and restore functions for variables. 

460 continue 

461 

462 if factory_data_list[0].name == trackable_utils.SERIALIZE_TO_TENSORS_NAME: 

463 # Trace Trackable save and restore functions. 

464 assert len(factory_data_list) == 1 

465 saveable_fn_map[obj] = {trackable_utils.SERIALIZE_TO_TENSORS_NAME: ( 

466 tracing_utils.trace_save_and_restore(obj))} 

467 else: 

468 # Trace deprecated SaveableObject save and restore functions. 

469 saveable_fn_map[obj] = ( 

470 trace_saveable_util.trace_save_restore_function_map( 

471 obj, factory_data_list)) 

472 return saveable_fn_map 

473 

474 

475def _tensor_dict_to_tensorinfo(tensor_dict): 

476 return { 

477 key: utils_impl.build_tensor_info_internal(value) 

478 for key, value in tensor_dict.items() 

479 } 

480 

481 

482def _to_safe_name_scope(signature_key, user_input_name): 

483 """Creates a sanitized name scope from user signature and input names. 

484 

485 Concatenates signature and input names, sanitizing as needed to be a valid 

486 scope name. 

487 

488 Args: 

489 signature_key: The user-provided key for the signature. 

490 user_input_name: The user-provided name for the input placeholder. 

491 

492 Returns: 

493 A name scope that is safe to be used in tf.name_scope(). 

494 """ 

495 name_scope = "{}_{}".format(signature_key, user_input_name) 

496 if re.match(r"^[A-Za-z0-9.][A-Za-z0-9_.\\-]*$", name_scope): 

497 return name_scope 

498 invalid_prefix_stripped = re.sub(r"^[^A-Za-z0-9.]*", "", name_scope) 

499 return re.sub(r"[^A-Za-z0-9_.\\-]", "_", invalid_prefix_stripped) 

500 

501 

502def _map_function_arguments_to_created_inputs( 

503 function_arguments, signature_key, function_name, defaults=None 

504): 

505 """Creates exterior placeholders in the exported graph for function arguments. 

506 

507 Functions have two types of inputs: tensors captured from the outside (eager) 

508 context, and arguments to the function which we expect to receive from the 

509 user at each call. `_map_captures_to_created_tensors` replaces 

510 captured tensors with stand-ins (typically these are resource dtype tensors 

511 associated with variables). `_map_function_inputs_to_created_inputs` runs over 

512 every argument, creating a new placeholder for each which will belong to the 

513 exported graph rather than the function body. 

514 

515 Args: 

516 function_arguments: A list of argument placeholders in the function body. 

517 signature_key: The name of the signature being exported, for error messages. 

518 function_name: The name of the function, for error messages. 

519 defaults: A dictionary mapping (signature_key, user_specified_name) to 

520 Tensor representing default values. 

521 

522 Returns: 

523 A tuple of (mapped_inputs, exterior_placeholders) 

524 mapped_inputs: A list with entries corresponding to `function_arguments` 

525 containing all of the inputs of the function gathered from the exported 

526 graph (both captured resources and arguments). 

527 exterior_argument_placeholders: A dictionary mapping from argument names 

528 to placeholders in the exported graph, containing the explicit arguments 

529 to the function which a user is expected to provide. 

530 

531 Raises: 

532 ValueError: If argument names are not unique. 

533 """ 

534 # `exterior_argument_placeholders` holds placeholders which are outside the 

535 # function body, directly contained in a MetaGraph of the SavedModel. The 

536 # function body itself contains nearly identical placeholders used when 

537 # running the function, but these exterior placeholders allow Session-based 

538 # APIs to call the function using feeds and fetches which name Tensors in the 

539 # MetaGraph. 

540 exterior_argument_placeholders = {} 

541 mapped_inputs = [] 

542 for placeholder in function_arguments: 

543 # `export_captures` contains an exhaustive set of captures, so if we don't 

544 # find the input there then we now know we have an argument. 

545 user_input_name = compat.as_str_any( 

546 placeholder.op.get_attr("_user_specified_name")) 

547 # If the internal placeholders for a function have names which were 

548 # uniquified by TensorFlow, then a single user-specified argument name 

549 # must refer to multiple Tensors. The resulting signatures would be 

550 # confusing to call. Instead, we throw an exception telling the user to 

551 # specify explicit names. 

552 if user_input_name != placeholder.op.name: 

553 # This should be unreachable, since concrete functions may not be 

554 # generated with non-unique argument names. 

555 raise ValueError( 

556 "Got non-flat/non-unique argument names for SavedModel signature " 

557 f"'{signature_key}': more than one argument to " 

558 f"'{compat.as_str_any(function_name)}' was named " 

559 f"'{user_input_name}'. " 

560 "Signatures have one Tensor per named input, so to have " 

561 "predictable names Python functions used to generate these " 

562 "signatures should avoid *args and Tensors in nested " 

563 "structures unless unique names are specified for each. Use " 

564 "tf.TensorSpec(..., name=...) to provide a name for a Tensor " 

565 "input.") 

566 default_value = defaults.get((signature_key, user_input_name)) 

567 if default_value is not None: 

568 placeholder_with_default = array_ops.placeholder_with_default( 

569 input=default_value.numpy(), 

570 shape=placeholder.shape, 

571 name=_to_safe_name_scope(signature_key, user_input_name), 

572 ) 

573 exterior_argument_placeholders[user_input_name] = placeholder_with_default 

574 mapped_inputs.append(placeholder_with_default) 

575 else: 

576 arg_placeholder = array_ops.placeholder( 

577 shape=placeholder.shape, 

578 dtype=placeholder.dtype, 

579 name=_to_safe_name_scope(signature_key, user_input_name), 

580 ) 

581 exterior_argument_placeholders[user_input_name] = arg_placeholder 

582 mapped_inputs.append(arg_placeholder) 

583 return mapped_inputs, exterior_argument_placeholders 

584 

585 

586def _generate_signatures(signature_functions, object_map, defaults=None): 

587 """Validates and calls `signature_functions` in the exported graph. 

588 

589 Args: 

590 signature_functions: A dictionary mapping string keys to concrete TensorFlow 

591 functions (e.g. from `signature_serialization.canonicalize_signatures`) 

592 which will be used to generate SignatureDefs. 

593 object_map: A dictionary that contains mappings from signature functions to 

594 concrete functions in the exported graph. 

595 defaults: A dictionary mapping (signature_key, user_specified_name) to 

596 Tensor representing default values. 

597 

598 Returns: 

599 Each function in the `signature_functions` dictionary is called with 

600 placeholder Tensors, generating a function call operation and output 

601 Tensors. The placeholder Tensors, the function call operation, and the 

602 output Tensors from the function call are part of the default Graph. 

603 

604 This function then returns a dictionary with the same structure as 

605 `signature_functions`, with the concrete functions replaced by SignatureDefs 

606 implicitly containing information about how to call each function from a 

607 TensorFlow 1.x Session / the C++ Loader API. These SignatureDefs reference 

608 the generated placeholders and Tensor outputs by name. 

609 

610 The caller is expected to include the default Graph set while calling this 

611 function as a MetaGraph in a SavedModel, including the returned 

612 SignatureDefs as part of that MetaGraph. 

613 """ 

614 signatures = {} 

615 for signature_key, function in sorted(signature_functions.items()): 

616 if function.graph.captures: 

617 argument_inputs = function.graph.inputs[:-len(function.graph.captures)] 

618 else: 

619 argument_inputs = function.graph.inputs 

620 mapped_inputs, exterior_argument_placeholders = ( 

621 _map_function_arguments_to_created_inputs( 

622 argument_inputs, signature_key, function.name, defaults 

623 ) 

624 ) 

625 kwarg_names = list( 

626 sorted( 

627 object_map[function].function.structured_input_signature[1].keys())) 

628 outputs = object_map[function](**{ 

629 kwarg_name: mapped_input 

630 for kwarg_name, mapped_input in zip(kwarg_names, mapped_inputs) 

631 }) 

632 signatures[signature_key] = signature_def_utils.build_signature_def( 

633 _tensor_dict_to_tensorinfo(exterior_argument_placeholders), 

634 _tensor_dict_to_tensorinfo(outputs), 

635 method_name=signature_constants.PREDICT_METHOD_NAME) 

636 return signatures 

637 

638 

639_AssetInfo = collections.namedtuple( 

640 "_AssetInfo", 

641 [ 

642 # List of AssetFileDef protocol buffers 

643 "asset_defs", 

644 # Map from asset variable resource Tensors to their init ops 

645 "asset_initializers_by_resource", 

646 # Map from base asset filenames to full paths 

647 "asset_filename_map", 

648 # Map from Asset to index of corresponding AssetFileDef 

649 "asset_index" 

650 ]) 

651 

652 

653def _add_asset_info(trackable_asset, asset_info, mapped_path_variable): 

654 """Add `trackable_asset` to `asset_info`.""" 

655 original_path_tensor = trackable_asset.asset_path 

656 original_path = tensor_util.constant_value(original_path_tensor) 

657 try: 

658 original_path = str(original_path.astype(str)) 

659 except AttributeError: 

660 # Already a string rather than a numpy array 

661 pass 

662 

663 path = builder_impl.get_asset_filename_to_add( 

664 asset_filepath=original_path, 

665 asset_filename_map=asset_info.asset_filename_map) 

666 asset_info.asset_filename_map[path] = original_path 

667 asset_def = meta_graph_pb2.AssetFileDef() 

668 asset_def.filename = path 

669 asset_def.tensor_info.name = mapped_path_variable.initial_value.name 

670 asset_info.asset_defs.append(asset_def) 

671 asset_info.asset_initializers_by_resource[original_path_tensor] = ( 

672 mapped_path_variable.initializer) 

673 asset_info.asset_index[trackable_asset] = len(asset_info.asset_defs) - 1 

674 

675 

676def _iterate_op_types(fn): 

677 """Iterates through each op in the function and returns the op type and op.""" 

678 if isinstance(fn, framework_fn._DefinedFunction): # pylint: disable=protected-access 

679 for node in fn.definition.node_def: 

680 op_type = node.attr["_gradient_op_type"].s 

681 if op_type: 

682 raise ValueError( 

683 "Unable to save gradient functions when exporting a " 

684 "_DefinedFunction (generally created through graph freezing utils " 

685 "or through V1 graph importers). Please save with " 

686 "`options=tf.SaveOptions(experimental_custom_gradients=False)`") 

687 else: 

688 for op in fn.graph.get_operations(): 

689 try: 

690 op_type = op.get_attr("_gradient_op_type") 

691 except ValueError: 

692 continue 

693 yield op_type, op 

694 

695 

696def _get_outer_most_capture(fn, capture, func_graph_map): 

697 """Tries to find the original captured tensor if capture more than once.""" 

698 outer_fn = fn 

699 while outer_fn is not None and not isinstance(capture, ops.EagerTensor): 

700 if capture.graph is not outer_fn.graph: 

701 outer_fn = func_graph_map.get(outer_fn.graph.outer_graph) 

702 else: 

703 try: 

704 capture_index = outer_fn.graph.internal_captures.index(capture) 

705 except ValueError: 

706 break # Capture is a tensor inside function, and not captured from 

707 # another external function 

708 capture = outer_fn.graph.external_captures[capture_index] 

709 outer_fn = func_graph_map.get(outer_fn.graph.outer_graph) 

710 return outer_fn, capture 

711 

712 

713def _trace_gradient_functions(graph, saveable_view): 

714 """Traces gradient functions and records them in the SaveableView.""" 

715 functions = list(graph._functions.values()) # pylint: disable=protected-access 

716 func_graph_map = {f.graph: f for f in functions if hasattr(f, "graph")} 

717 seen_op_types = set() 

718 

719 for fn in functions: 

720 for op_type, op in _iterate_op_types(fn): 

721 if op_type in seen_op_types: 

722 continue 

723 seen_op_types.add(op_type) 

724 

725 try: 

726 custom_gradient = ops.gradient_registry.lookup(op_type) 

727 except LookupError: 

728 continue 

729 

730 try: 

731 grad_fn = ( 

732 def_function.function(custom_gradient).get_concrete_function( 

733 None, *op.inputs)) 

734 except Exception as exc: 

735 traceback.print_exc() 

736 raise ValueError( 

737 "Error when tracing gradients for SavedModel.\n\n" 

738 "Check the error log to see the error that was raised when " 

739 "converting a gradient function to a concrete function. You may " 

740 "need to update the custom gradient, or disable saving gradients " 

741 "with the option " 

742 "tf.saved_model.SaveOptions(experimental_custom_gradients=False)" 

743 f".\n\tProblematic op name: {op.name}\n\tGradient inputs: " 

744 f"{op.inputs}") from exc 

745 

746 with graph.as_default(): 

747 # The gradient function will capture all intermediate values. These 

748 # captures be serialized so that they can be re-bound to the function 

749 # when loading. 

750 bad_captures = [] 

751 for capture in grad_fn.captured_inputs: 

752 if capture.dtype in _UNCOPIABLE_DTYPES: 

753 continue 

754 # Tries to find the outermost capture in case the tensor is a constant 

755 # or not actually captured in the current function (this could happen 

756 # if the function is a while loop body, in which case the captured 

757 # input is not the internal captured tensor). 

758 outer_fn, outer_capture = _get_outer_most_capture( 

759 fn, capture, func_graph_map 

760 ) 

761 if outer_fn is None or isinstance(outer_capture, ops.EagerTensor): 

762 if outer_capture not in saveable_view.captured_tensor_node_ids: 

763 raise ValueError( 

764 f"Found invalid capture {outer_capture} when " 

765 "saving custom gradients." 

766 ) 

767 saveable_view.captured_tensor_node_ids[capture] = ( 

768 saveable_view.captured_tensor_node_ids[outer_capture] 

769 ) 

770 elif outer_capture.graph is outer_fn.graph: 

771 capture_name = outer_capture.name 

772 # It's possible for AtomicFunctions to save different names 

773 # for input tensors when serialized to FunctionDef (all 

774 # non-alphanumeric characters are converted to '_'). 

775 if isinstance(outer_fn, defun.AtomicFunction): # pylint:disable=protected-access 

776 try: 

777 arg_index = outer_fn.graph.inputs.index(outer_capture) 

778 capture_name = ( 

779 outer_fn.cached_definition.signature.input_arg[ 

780 arg_index 

781 ].name 

782 + ":0" 

783 ) 

784 except ValueError: 

785 pass 

786 

787 node = _CapturedTensor(capture_name, outer_fn.name) 

788 saveable_view.add_capture_and_node(capture, node) 

789 else: 

790 bad_captures.append(capture.name) 

791 if not bad_captures: 

792 grad_fn.add_to_graph(graph) 

793 else: 

794 raise ValueError( 

795 f"Cannot save custom gradient {op_type} called in function {fn} " 

796 "because SavedModel is unable to serialize the captured " 

797 f"inputs: {bad_captures}" 

798 ) 

799 

800 saveable_view.gradient_functions.append(grad_fn) 

801 func_graph_map[grad_fn.graph] = grad_fn 

802 

803 grad_def = function_pb2.RegisteredGradient() 

804 grad_def.gradient_func = grad_fn.name 

805 grad_def.registered_op_type = op_type 

806 saveable_view.gradient_defs.append(grad_def) 

807 

808 

809def _fill_meta_graph_def( 

810 meta_graph_def, 

811 saveable_view, 

812 signature_functions, 

813 namespace_whitelist, 

814 save_custom_gradients, 

815 defaults=None, 

816): 

817 """Generates a MetaGraph which calls `signature_functions`. 

818 

819 Args: 

820 meta_graph_def: The MetaGraphDef proto to fill. 

821 saveable_view: The _SaveableView being exported. 

822 signature_functions: A dictionary mapping signature keys to concrete 

823 functions containing signatures to add to the MetaGraph. 

824 namespace_whitelist: List of strings containing whitelisted op namespaces. 

825 save_custom_gradients: Whether to save custom gradients. 

826 defaults: A dictionary mapping (signature_key, user_specified_name) to 

827 Tensor representing default values. 

828 

829 Returns: 

830 A tuple of (_AssetInfo, Graph) containing the captured assets and 

831 exported Graph generated from tracing the saveable_view. 

832 """ 

833 # List objects from the eager context to make sure Optimizers give us the 

834 # right Graph-dependent variables. 

835 resource_initializers = saveable_view.get_concrete_resource_initializers() 

836 exported_graph = ops.Graph() 

837 resource_initializer_ops = [] 

838 with exported_graph.as_default(): 

839 object_map, tensor_map, asset_info = saveable_view.map_resources() 

840 signatures = _generate_signatures(signature_functions, object_map, defaults) 

841 if save_custom_gradients: 

842 # Custom gradients functions must be traced in the same context as the 

843 # when they are registered. 

844 _trace_gradient_functions(exported_graph, saveable_view) 

845 with exported_graph.as_default(): 

846 # Create initializers for assets and resources. 

847 for resource_initializer_function in resource_initializers: 

848 asset_dependencies = [] 

849 for capture in resource_initializer_function.graph.external_captures: 

850 asset_initializer = asset_info.asset_initializers_by_resource.get( 

851 capture, None) 

852 if asset_initializer is not None: 

853 asset_dependencies.append(asset_initializer) 

854 with ops.control_dependencies(asset_dependencies): 

855 mapped_initializer = object_map[resource_initializer_function] 

856 resource_initializer_ops.append(mapped_initializer()) 

857 resource_initializer_ops.extend( 

858 asset_info.asset_initializers_by_resource.values()) 

859 with ops.control_dependencies(resource_initializer_ops): 

860 init_op = control_flow_ops.no_op() 

861 # Add the same op to the main_op collection and to the init_op 

862 # signature. The collection is for compatibility with older loader APIs; 

863 # only one will be executed. 

864 meta_graph_def.collection_def[constants.MAIN_OP_KEY].node_list.value.append( 

865 init_op.name) 

866 meta_graph_def.signature_def[constants.INIT_OP_SIGNATURE_KEY].CopyFrom( 

867 signature_def_utils.op_signature_def(init_op, 

868 constants.INIT_OP_SIGNATURE_KEY)) 

869 

870 # Saving an object-based checkpoint again gathers variables. We need to do the 

871 # gathering from the eager context so Optimizers save the right set of 

872 # variables, but want any operations associated with the save/restore to be in 

873 # the exported graph (thus the `to_graph` argument). 

874 def call_with_mapped_captures(function, args): 

875 if function in object_map: 

876 return object_map[function](*args) 

877 # Registered saver/restore functions do not appear in `object_map`, because 

878 # they are not in the object graph. 

879 return saved_model_exported_concrete.ExportedConcreteFunction( 

880 function, tensor_map)(*args) 

881 

882 for obj in object_map.values():