TensorFlow核心组件系列之Graph的底层机制探索

class Graph(object):
       def __init__(self):
         self._lock = threading.Lock()
         self._nodes_by_id = dict()    # GUARDED_BY(self._lock)
         self._next_id_counter = 0     # GUARDED_BY(self._lock)
         self._nodes_by_name = dict()  # GUARDED_BY(self._lock)
         self._registered_ops = op_def_registry.get_registered_ops()

        def _add_op(self, op):
                self._check_not_finalized()
                if not isinstance(op, (Tensor, Operation)):
                  raise TypeError("op must be a Tensor or Operation: %s" % op)
                with self._lock:
                  # pylint: disable=protected-access
                  if op._id in self._nodes_by_id:
                    raise ValueError("cannot add an op with id %d as it already "
                                     "exists in the graph" % op._id)
                  if op.name in self._nodes_by_name:
                    raise ValueError("cannot add op with name %s as that name "
                                     "is already used" % op.name)
                  self._nodes_by_id[op._id] = op
                  self._nodes_by_name[op.name] = op
                  self._version = max(self._version, op._id)

         def add_to_collection(name, value):
                      get_default_graph().add_to_collection(name, value)

         # 替换线程默认图
               def as_default(self):
                   return _default_graph_stack.get_controller(self)

               # 栈式管理，push pop
               @tf_contextlib.contextmanager
               def get_controller(self, default):
                    try:
                      context.context_stack.push(default.building_function, default.as_default)
                    finally:
                      context.context_stack.pop()

         def create_op(
                  self,
                  op_type,
                  inputs,
                  dtypes=None,  # pylint: disable=redefined-outer-name
                  input_types=None,
                  name=None,
                  attrs=None,
                  op_def=None,
                  compute_shapes=True,
                  compute_device=True):
                        for idx, a in enumerate(inputs):
                          if not isinstance(a, Tensor):
                            raise TypeError("Input #%d is not a tensor: %s" % (idx, a))
                        return self._create_op_internal(op_type, inputs, dtypes, input_types, name,
                                                        attrs, op_def, compute_device)

print tf.get_default_graph()

with tf.Graph().as_default() as g:
    print tf.get_default_graph()

print tf.get_default_graph()

<tensorflow.python.framework.ops.Graph object at 0x106329fd0>
<tensorflow.python.framework.ops.Graph object at 0x18205cc0d0>
<tensorflow.python.framework.ops.Graph object at 0x10d025fd0>

class GraphKeys(object):
    GLOBAL_VARIABLES = "variables"
  # Key to collect local variables that are local to the machine and are not
  # saved/restored.
  LOCAL_VARIABLES = "local_variables"
  # optimizers.
  TRAINABLE_VARIABLES = "trainable_variables"
  SAVERS = "savers"
  # Key to collect weights
  WEIGHTS = "weights"
  # Key to collect biases
  BIASES = "biases"
  # Key to collect activations
  ACTIVATIONS = "activations"
  # Key to collect update_ops
  UPDATE_OPS = "update_ops"
  # Key to collect losses
  LOSSES = "losses"
  ...

  # Key to indicate various ops.
  INIT_OP = "init_op"
  LOCAL_INIT_OP = "local_init_op"
  SUMMARY_OP = "summary_op"
  GLOBAL_STEP = "global_step"

  # Used to count the number of evaluations performed during a single evaluation
  # run.
  EVAL_STEP = "eval_step"
  TRAIN_OP = "train_op"

  # Key for control flow context.
  COND_CONTEXT = "cond_context"
  WHILE_CONTEXT = "while_context"

  # Used to store v2 summary names.
  _SUMMARY_COLLECTION = "_SUMMARY_V2"

  # List of all collections that keep track of variables.
  _VARIABLE_COLLECTIONS = [
      GLOBAL_VARIABLES,
      LOCAL_VARIABLES,
      METRIC_VARIABLES,
      MODEL_VARIABLES,
      TRAINABLE_VARIABLES,
      MOVING_AVERAGE_VARIABLES,
      CONCATENATED_VARIABLES,
      TRAINABLE_RESOURCE_VARIABLES,
  ]

# 用户要快速的检索某类变量可以通过这样的语句
all_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)

class Graph {
     private:
      // 所有已知的op计算函数的注册表
      FunctionLibraryDefinition ops_;

      // GraphDef版本号
      const std::unique_ptr<VersionDef> versions_;

      // 节点node列表，通过id来访问
      std::vector<Node*> nodes_;

      // node个数
      int64 num_nodes_ = 0;

      // 边edge列表，通过id来访问
      std::vector<Edge*> edges_;

      // graph中非空edge的数目
      int num_edges_ = 0;

      // 已分配了内存，但还没使用的node和edge
      std::vector<Node*> free_nodes_;
      std::vector<Edge*> free_edges_;

     const Edge* Graph::AddEdge(Node* source, int x, Node* dest, int y) {
       auto e = AllocEdge();
       e->src_ = source;
       e->dst_ = dest;
       e->src_output_ = x;
       e->dst_input_ = y;
       CHECK(source->out_edges_.insert(e).second);
       CHECK(dest->in_edges_.insert(e).second);
       edges_.push_back(e);
       edge_set_.insert(e);
       return e;
        }
 }