def _process_asset(trackable_asset, asset_info, resource_map): """Add `trackable_asset` to `asset_info` and `resource_map`.""" original_variable = trackable_asset.asset_path with context.eager_mode(): original_path = original_variable.numpy() path = builder_impl.get_asset_filename_to_add( asset_filepath=original_path, asset_filename_map=asset_info.asset_filename_map) # TODO(andresp): Instead of mapping 1-1 between trackable asset # and asset in the graph def consider deduping the assets that # point to the same file. asset_path_initializer = array_ops.placeholder( shape=original_variable.shape, dtype=dtypes.string, name="asset_path_initializer") asset_variable = resource_variable_ops.ResourceVariable( asset_path_initializer) asset_info.asset_filename_map[path] = original_path asset_def = meta_graph_pb2.AssetFileDef() asset_def.filename = path asset_def.tensor_info.name = asset_path_initializer.name asset_info.asset_defs.append(asset_def) asset_info.asset_initializers_by_resource[original_variable.handle] = ( asset_variable.initializer) asset_info.asset_index[trackable_asset] = len(asset_info.asset_defs) - 1 resource_map[original_variable.handle] = asset_variable.handle
def _process_asset(trackable_asset, asset_info, resource_map): """Add `trackable_asset` to `asset_info` and `resource_map`.""" original_path_tensor = trackable_asset.asset_path original_path = tensor_util.constant_value(original_path_tensor) try: original_path = str(original_path.astype(str)) except AttributeError: # Already a string rather than a numpy array pass path = builder_impl.get_asset_filename_to_add( asset_filepath=original_path, asset_filename_map=asset_info.asset_filename_map) # TODO(andresp): Instead of mapping 1-1 between trackable asset # and asset in the graph def consider deduping the assets that # point to the same file. asset_path_initializer = array_ops.placeholder( shape=original_path_tensor.shape, dtype=dtypes.string, name="asset_path_initializer") asset_variable = resource_variable_ops.ResourceVariable( asset_path_initializer) asset_info.asset_filename_map[path] = original_path asset_def = meta_graph_pb2.AssetFileDef() asset_def.filename = path asset_def.tensor_info.name = asset_path_initializer.name asset_info.asset_defs.append(asset_def) asset_info.asset_initializers_by_resource[original_path_tensor] = ( asset_variable.initializer) asset_info.asset_index[trackable_asset] = len(asset_info.asset_defs) - 1 resource_map[original_path_tensor] = asset_variable
def testCreationOfAssetsKeyCollectionIsDeterministic(self): tmp_asset_dir = os.path.join(self.get_temp_dir(), "assets") tf_v1.gfile.MakeDirs(tmp_asset_dir) filenames = [ os.path.join(tmp_asset_dir, "file%d.txt" % n) for n in range(10) ] for filename in filenames: _write_string_to_file(filename, "I am file %s" % filename) with tf.Graph().as_default() as graph: assets = [ tf.constant(f, name=os.path.basename(f)) for f in filenames ] for asset in assets: graph.add_to_collection(tf_v1.GraphKeys.ASSET_FILEPATHS, asset) saved_model_lib.add_signature("default", {}, {"default": assets[0]}) handler = saved_model_lib.SavedModelHandler() handler.add_graph_copy(graph) saved_model_proto = copy.deepcopy(handler._proto) export_dir = os.path.join(self.get_temp_dir(), "assets_key_test") saved_model_lib._make_assets_key_collection(saved_model_proto, export_dir) meta_graph = list(saved_model_proto.meta_graphs)[0] asset_tensor_names = [] for asset_any_proto in meta_graph.collection_def[ tf_v1.saved_model.constants.ASSETS_KEY].any_list.value: asset_proto = meta_graph_pb2.AssetFileDef() asset_any_proto.Unpack(asset_proto) asset_tensor_names.append(asset_proto.tensor_info.name) self.assertEqual(asset_tensor_names, sorted(asset_tensor_names))
def _merge_assets_key_collection(saved_model_proto, path): """Merges the ASSETS_KEY collection into the GraphDefs in saved_model_proto. Removes the ASSETS_KEY collection from the GraphDefs in the SavedModel and modifies nodes with the assets filenames to point to the assets in `path`. After this transformation, the SavedModel GraphDefs can be used without feeding asset tensors. Args: saved_model_proto: SavedModel proto to be modified. path: path where the SavedModel is being loaded from. """ for meta_graph in saved_model_proto.meta_graphs: node_asset_map = {} if tf_v1.saved_model.constants.ASSETS_KEY in meta_graph.collection_def: assets_any_proto = meta_graph.collection_def[ tf_v1.saved_model.constants.ASSETS_KEY].any_list.value for asset_any_proto in assets_any_proto: asset_proto = meta_graph_pb2.AssetFileDef() asset_any_proto.Unpack(asset_proto) asset_filename = _get_asset_filename(path, asset_proto.filename) node_asset_map[_get_node_name_from_tensor( asset_proto.tensor_info.name)] = asset_filename del meta_graph.collection_def[ tf_v1.saved_model.constants.ASSETS_KEY] for node in meta_graph.graph_def.node: asset_filepath = node_asset_map.get(node.name) if asset_filepath: _check_asset_node_def(node) node.attr["value"].tensor.string_val[0] = asset_filepath
def _get_asset_tensors(export_dir, meta_graph_def_to_load): """Gets the asset tensors, if defined in the meta graph def to load. Args: export_dir: Directory where the SavedModel is located. meta_graph_def_to_load: The meta graph def from the SavedModel to be loaded. Returns: A dictionary of asset tensors, keyed by the name of the asset tensor. The value in the map corresponds to the absolute path of the asset file. """ # Collection-def that may contain the assets key. collection_def = meta_graph_def_to_load.collection_def asset_tensor_dict = {} if constants.ASSETS_KEY in collection_def: # Location of the assets for SavedModel. assets_directory = os.path.join( compat.as_bytes(export_dir), compat.as_bytes(constants.ASSETS_DIRECTORY)) assets_any_proto = collection_def[constants.ASSETS_KEY].any_list.value # Process each asset and add it to the asset tensor dictionary. for asset_any_proto in assets_any_proto: asset_proto = meta_graph_pb2.AssetFileDef() asset_any_proto.Unpack(asset_proto) asset_tensor_dict[asset_proto.tensor_info.name] = os.path.join( compat.as_bytes(assets_directory), compat.as_bytes(asset_proto.filename)) return asset_tensor_dict
def get_asset_tensors(export_dir, meta_graph_def_to_load, import_scope=None): """Gets the asset tensors, if defined in the meta graph def to load. Args: export_dir: Directory where the SavedModel is located. meta_graph_def_to_load: The meta graph def from the SavedModel to be loaded. import_scope: Optional `string` -- if specified, prepend this followed by '/' to all returned asset tensor names. Returns: A dictionary of asset tensors, keyed by the name of the asset tensor. The value in the map corresponds to the absolute path of the asset file. """ # Collection-def that may contain the assets key. collection_def = meta_graph_def_to_load.collection_def asset_tensor_dict = {} if constants.ASSETS_KEY in collection_def: # Location of the assets for SavedModel. assets_directory = os.path.join( compat.as_bytes(export_dir), compat.as_bytes(constants.ASSETS_DIRECTORY)) assets_any_proto = collection_def[constants.ASSETS_KEY].any_list.value # Process each asset and add it to the asset tensor dictionary. for asset_any_proto in assets_any_proto: asset_proto = meta_graph_pb2.AssetFileDef() asset_any_proto.Unpack(asset_proto) tensor_name = asset_proto.tensor_info.name if import_scope: tensor_name = "%s/%s" % (import_scope, tensor_name) asset_tensor_dict[tensor_name] = os.path.join( compat.as_bytes(assets_directory), compat.as_bytes(asset_proto.filename)) return asset_tensor_dict
def testAssets(self): export_dir = os.path.join(compat.as_bytes(tf.test.get_temp_dir()), compat.as_bytes("with-assets")) builder = saved_model_builder.SavedModelBuilder(export_dir) with self.test_session(graph=tf.Graph()) as sess: v = tf.Variable(42, name="v") sess.run(tf.initialize_all_variables()) self.assertEqual(42, v.eval()) # Build an asset collection. asset_filepath = os.path.join( compat.as_bytes(tf.test.get_temp_dir()), compat.as_bytes("hello42.txt")) file_io.write_string_to_file(asset_filepath, "foo bar baz") asset_file_tensor = tf.constant(asset_filepath, name="asset_file_tensor") tf.add_to_collection(tf.GraphKeys.ASSET_FILEPATHS, asset_file_tensor) ignored_filepath = os.path.join( compat.as_bytes(tf.test.get_temp_dir()), compat.as_bytes("ignored.txt")) file_io.write_string_to_file(ignored_filepath, "will be ignored") asset_collection = tf.get_collection(tf.GraphKeys.ASSET_FILEPATHS) builder.add_meta_graph_and_variables( sess, ["foo"], assets_collection=asset_collection) # Save the SavedModel to disk. builder.save() with self.test_session(graph=tf.Graph()) as sess: foo_graph = loader.load(sess, ["foo"], export_dir) # Validate the assets. collection_def = foo_graph.collection_def assets_any = collection_def[constants.ASSETS_KEY].any_list.value self.assertEqual(len(assets_any), 1) asset = meta_graph_pb2.AssetFileDef() assets_any[0].Unpack(asset) assets_path = os.path.join( compat.as_bytes(export_dir), compat.as_bytes(constants.ASSETS_DIRECTORY), compat.as_bytes("hello42.txt")) asset_contents = file_io.read_file_to_string(assets_path) self.assertEqual("foo bar baz", compat.as_text(asset_contents)) self.assertEqual("hello42.txt", asset.filename) self.assertEqual("asset_file_tensor:0", asset.tensor_info.name) ignored_asset_path = os.path.join( compat.as_bytes(export_dir), compat.as_bytes(constants.ASSETS_DIRECTORY), compat.as_bytes("ignored.txt")) self.assertFalse(file_io.file_exists(ignored_asset_path))
def _add_asset_to_collection(asset_filename, asset_tensor): """Builds an asset proto and adds it to the asset collection of the graph. Args: asset_filename: The filename of the asset to be added. asset_tensor: The asset tensor used to populate the tensor info of the asset proto. """ asset_proto = meta_graph_pb2.AssetFileDef() asset_proto.filename = asset_filename asset_proto.tensor_info.name = asset_tensor.name asset_any_proto = Any() asset_any_proto.Pack(asset_proto) ops.add_to_collection(constants.ASSETS_KEY, asset_any_proto)
def _validate_asset_collection(self, export_dir, graph_collection_def, expected_asset_file_name, expected_asset_file_contents, expected_asset_tensor_name): assets_any = graph_collection_def[constants.ASSETS_KEY].any_list.value asset = meta_graph_pb2.AssetFileDef() assets_any[0].Unpack(asset) assets_path = os.path.join(compat.as_bytes(export_dir), compat.as_bytes(constants.ASSETS_DIRECTORY), compat.as_bytes(expected_asset_file_name)) actual_asset_contents = file_io.read_file_to_string(assets_path) self.assertEqual(expected_asset_file_contents, compat.as_text(actual_asset_contents)) self.assertEqual(expected_asset_file_name, asset.filename) self.assertEqual(expected_asset_tensor_name, asset.tensor_info.name)
def _process_asset(trackable_asset, asset_info, resource_map): """Add `trackable_asset` to `asset_info` and `resource_map`.""" original_variable = trackable_asset.asset_path with context.eager_mode(): original_path = original_variable.numpy() path = builder_impl.get_asset_filename_to_add( asset_filepath=original_path, asset_filename_map=asset_info.asset_filename_map) asset_variable = asset_info.asset_filename_map.get(path, None) if asset_variable is None: asset_path_initializer = array_ops.placeholder( shape=original_variable.shape, dtype=dtypes.string, name="asset_path_initializer") asset_variable = resource_variable_ops.ResourceVariable( asset_path_initializer) asset_info.asset_filename_map[path] = original_path asset_def = meta_graph_pb2.AssetFileDef() asset_def.filename = path asset_def.tensor_info.name = asset_path_initializer.name asset_info.asset_defs.append(asset_def) asset_info.asset_initializers_by_resource[original_variable.handle] = ( asset_variable.initializer) resource_map[original_variable.handle] = asset_variable.handle
def _make_assets_key_collection(saved_model_proto, export_path): """Creates an ASSETS_KEY collection in the GraphDefs in saved_model_proto. Adds an ASSETS_KEY collection to the GraphDefs in the SavedModel and returns a map from original asset filename to filename when exporting the SavedModel to `export_path`. This is roughly the inverse operation of `_merge_assets_key_collection`. Args: saved_model_proto: SavedModel proto to be modified. export_path: string with path where the saved_model_proto will be exported. Returns: A map from original asset filename to asset filename when exporting the SavedModel to path. Raises: ValueError: on unsuported/unexpected SavedModel. """ asset_filenames = {} used_asset_filenames = set() def _make_asset_filename(original_filename): """Returns the asset filename to use for the file.""" if original_filename in asset_filenames: return asset_filenames[original_filename] basename = os.path.basename(original_filename) suggestion = basename index = 0 while suggestion in used_asset_filenames: suggestion = tf.compat.as_bytes(basename) + tf.compat.as_bytes( str(index)) index += 1 asset_filenames[original_filename] = suggestion used_asset_filenames.add(suggestion) return suggestion for meta_graph in saved_model_proto.meta_graphs: collection_def = meta_graph.collection_def.get( tf_v1.GraphKeys.ASSET_FILEPATHS) if collection_def is None: continue if collection_def.WhichOneof("kind") != "node_list": raise ValueError( "MetaGraph collection ASSET_FILEPATHS is not a list of tensors." ) for tensor in collection_def.node_list.value: if not tensor.endswith(":0"): raise ValueError( "Unexpected tensor in ASSET_FILEPATHS collection.") asset_nodes = set([ _get_node_name_from_tensor(tensor) for tensor in collection_def.node_list.value ]) tensor_filename_map = {} for node in meta_graph.graph_def.node: if node.name in asset_nodes: _check_asset_node_def(node) filename = node.attr["value"].tensor.string_val[0] tensor_filename_map[node.name + ":0"] = filename logging.debug("Found asset node %s pointing to %s", node.name, filename) # Clear value to avoid leaking the original path. node.attr["value"].tensor.string_val[0] = ( tf.compat.as_bytes("SAVEDMODEL-ASSET")) if tensor_filename_map: assets_key_collection = meta_graph.collection_def[ tf_v1.saved_model.constants.ASSETS_KEY] for tensor, filename in sorted(tensor_filename_map.items()): asset_proto = meta_graph_pb2.AssetFileDef() asset_proto.filename = _make_asset_filename(filename) asset_proto.tensor_info.name = tensor assets_key_collection.any_list.value.add().Pack(asset_proto) return { original_filename: _get_asset_filename(export_path, asset_filename) for original_filename, asset_filename in asset_filenames.items() }
def _make_asset_file_def_any(node_name): asset_file_def = meta_graph_pb2.AssetFileDef() asset_file_def.tensor_info.name = node_name any_message = Any() any_message.Pack(asset_file_def) return any_message
def _build_model(self): """ load graph from model_path and create session for this graph """ model_path = self._model_path self._graph = tf.Graph() gpu_options = tf.GPUOptions(allow_growth=True) session_config = tf.ConfigProto(gpu_options=gpu_options, allow_soft_placement=True, log_device_placement=False) self._session = tf.Session(config=session_config, graph=self._graph) with self._graph.as_default(): with self._session.as_default(): # load model _, ext = os.path.splitext(model_path) tf.logging.info('loading model from %s' % model_path) if tf.gfile.IsDirectory(model_path): model_path = self.search_pb(model_path) logging.info('model find in %s' % model_path) assert tf.saved_model.loader.maybe_saved_model_directory(model_path), \ 'saved model does not exists in %s' % model_path self._is_saved_model = True meta_graph_def = tf.saved_model.loader.load( self._session, [tf.saved_model.tag_constants.SERVING], model_path) # parse signature signature_def = meta_graph_def.signature_def[ signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] inputs = signature_def.inputs for name, tensor in inputs.items(): logging.info('Load input binding: %s -> %s' % (name, tensor.name)) self._inputs_map[ name] = self._graph.get_tensor_by_name(tensor.name) outputs = signature_def.outputs for name, tensor in outputs.items(): logging.info('Load output binding: %s -> %s' % (name, tensor.name)) self._outputs_map[ name] = self._graph.get_tensor_by_name(tensor.name) # get assets self._assets = {} asset_files = tf.get_collection(constants.ASSETS_KEY) for any_proto in asset_files: asset_file = meta_graph_pb2.AssetFileDef() any_proto.Unpack(asset_file) type_name = asset_file.tensor_info.name.split(':')[0] asset_path = os.path.join(model_path, constants.ASSETS_DIRECTORY, asset_file.filename) assert tf.gfile.Exists( asset_path ), '%s is missing in saved model' % asset_path self._assets[type_name] = asset_path logging.info(self._assets) # get export config self._export_config = {} self._use_bgr = False export_config_collection = tf.get_collection( 'EV_EXPORT_CONFIG') if len(export_config_collection) > 0: self._export_config = json.loads( export_config_collection[0]) logging.info('load export config info %s' % export_config_collection[0]) self._use_bgr = self._export_config.get( 'color_format', 'rgb').lower() == 'bgr' if self._use_bgr: logging.info( 'prediction will use image in bgr order') else: logging.info( 'prediction will use image in rgb order') else: raise ValueError('saved model is not found in %s' % self._model_path)