def _find_ends_and_bottleneck_tensors()

in tcav/model.py [0:0]

• 11 lines of code
• 5 McCabe index (conditional complexity)

  def _find_ends_and_bottleneck_tensors(self, node_dict):
    """ Find tensors from the graph by their names.

    Depending on how the model is loaded, tensors in the graph
    may or may not have 'import/' prefix added to every tensor name.
    This is true even if the tensors already have 'import/' prefix.
    The 'ends' and 'bottlenecks_tensors' dictionary should map to tensors
    with the according name.
    """
    self.bottlenecks_tensors = {}
    self.ends = {}
    for k, v in six.iteritems(node_dict):
      if self.import_prefix:
        v = 'import/' + v
      tensor = self.sess.graph.get_operation_by_name(v.strip(':0')).outputs[0]
      if k == 'input' or k == 'prediction':
        self.ends[k] = tensor
      else:
        self.bottlenecks_tensors[k] = tensor