Python tf_optimizer.OptimizeGraph方法代码示例

# 需要导入模块: from tensorflow.python.grappler import tf_optimizer [as 别名]
# 或者: from tensorflow.python.grappler.tf_optimizer import OptimizeGraph [as 别名]
def auto_parallel(metagraph, model):
  from tensorflow.python.grappler import tf_optimizer
  rewriter_config = rewriter_config_pb2.RewriterConfig()
  rewriter_config.optimizers.append("autoparallel")
  rewriter_config.auto_parallel.enable = True
  rewriter_config.auto_parallel.num_replicas = FLAGS.num_gpus
  optimized_graph = tf_optimizer.OptimizeGraph(rewriter_config, metagraph)
  metagraph.graph_def.CopyFrom(optimized_graph)
  UpdateCollection(metagraph, model) 

# 需要导入模块: from tensorflow.python.grappler import tf_optimizer [as 别名]
# 或者: from tensorflow.python.grappler.tf_optimizer import OptimizeGraph [as 别名]
def optimize_graph_spec(graph_spec_obj, config_proto):
  """Applies Grappler with given options to a `graph_spec.GraphSpec`.

  For more information on Grappler, see
  https://www.tensorflow.org/guide/graph_optimization

  Args:
    graph_spec_obj: Instance of `graph_spec.GraphSpec` representing the
      TensorFlow computation to optimize.
    config_proto: Instance of `tf.compat.v1.ConfigProto` specifying optimization
      options for Grappler.

  Returns:
    An instance of `graph_spec_obj` which has been passed through Grappler and
    optimized if possible.
  """
  meta_graph_def = graph_spec_obj.to_meta_graph_def()

  try:
    # Grappler raises if it fails to find feeds and fetches, but can handle
    # *some* no-arg graphs, so we try/except here.
    optimized_graph_def = tf_optimizer.OptimizeGraph(config_proto,
                                                     meta_graph_def)
  except ValueError as error:
    logging.info(
        'Grappler has raised the error %s; falling back to using '
        'non-optimized graph.', error)
    optimized_graph_def = graph_spec_obj.graph_def

  return graph_spec.GraphSpec(
      optimized_graph_def,
      init_op=graph_spec_obj.init_op,
      in_names=graph_spec_obj.in_names,
      out_names=graph_spec_obj.out_names) 

# 需要导入模块: from tensorflow.python.grappler import tf_optimizer [as 别名]
# 或者: from tensorflow.python.grappler.tf_optimizer import OptimizeGraph [as 别名]
def _run_tf_optimizer(config: ConfigProto,
                      graph: tf.Graph,
                      signature_def: SignatureDef) -> GraphDef:
    """Run the TF optimizer ("grappler") on a graph"""
    graph_def = graph.as_graph_def()
    meta_graph = export_meta_graph(graph_def=graph_def, graph=graph)
    meta_graph.signature_def['not_used_key'].CopyFrom(signature_def)
    return tf_optimizer.OptimizeGraph(config, meta_graph) 

# 需要导入模块: from tensorflow.python.grappler import tf_optimizer [as 别名]
# 或者: from tensorflow.python.grappler.tf_optimizer import OptimizeGraph [as 别名]
def _run_inline_graph_optimization(func, lower_control_flow):
  """Apply function inline optimization to the graph.

  Returns the GraphDef after Grappler's function inlining optimization is
  applied. This optimization does not work on models with control flow.

  Args:
    func: ConcreteFunction.
    lower_control_flow: Boolean indicating whether or not to lower control flow
      ops such as If and While. (default True)

  Returns:
    GraphDef
  """
  graph_def = func.graph.as_graph_def()
  if not lower_control_flow:
    graph_def = disable_lower_using_switch_merge(graph_def)

  # In some cases, a secondary implementation of the function (e.g. for GPU) is
  # written to the "api_implements" attribute. (e.g. `tf.keras.layers.LSTM` in
  # TF2 produces a CuDNN-based RNN for GPU).
  # This function suppose to inline all functions calls, but "api_implements"
  # prevents this from happening. Removing the attribute solves the problem.
  # To learn more about "api_implements", see:
  #   tensorflow/core/grappler/optimizers/implementation_selector.h
  for function in graph_def.library.function:
    if "api_implements" in function.attr:
      del function.attr["api_implements"]

  meta_graph = export_meta_graph(graph_def=graph_def, graph=func.graph)

  # Clear the initializer_name for the variables collections, since they are not
  # needed after saved to saved_model.
  for name in [
      "variables", "model_variables", "trainable_variables", "local_variables"
  ]:
    raw_list = []
    for raw in meta_graph.collection_def["variables"].bytes_list.value:
      variable = variable_pb2.VariableDef()
      variable.ParseFromString(raw)
      variable.ClearField("initializer_name")
      raw_list.append(variable.SerializeToString())
    meta_graph.collection_def[name].bytes_list.value[:] = raw_list

  # Add a collection 'train_op' so that Grappler knows the outputs.
  fetch_collection = meta_graph_pb2.CollectionDef()
  for array in func.inputs + func.outputs:
    fetch_collection.node_list.value.append(array.name)
  meta_graph.collection_def["train_op"].CopyFrom(fetch_collection)

  # Initialize RewriterConfig with everything disabled except function inlining.
  config = config_pb2.ConfigProto()
  rewrite_options = config.graph_options.rewrite_options
  rewrite_options.min_graph_nodes = -1  # do not skip small graphs
  rewrite_options.optimizers.append("function")
  return tf_optimizer.OptimizeGraph(config, meta_graph)