TensorFlow Lite の FlexDelegate (その2)

@Vengineerの戯言 : Twitter
SystemVerilogの世界へようこそ、すべては、SystemC v0.9公開から始まった

昨日のTensorFlow Lite の FlexDelegate (その1)の続き、

Subgraph::ReplaceNodeSubsetsWithDelegateKernels メソッド

TfLiteStatus Subgraph::ReplaceNodeSubsetsWithDelegateKernels(
    TfLiteRegistration registration, const TfLiteIntArray* nodes_to_replace,
    TfLiteDelegate* delegate) {
  // Annotate the registration as DELEGATE op.
  registration.builtin_code = BuiltinOperator_DELEGATE;

  // Analyze the graph to find all independent node_subsets that are either
  // fully not-this-delegate or this-delegate computation.
  InterpreterInfo info(this);
  std::vector<NodeSubset> node_subsets;
  PartitionGraphIntoIndependentNodeSubsets(&info, nodes_to_replace,
                                           &node_subsets);

PartitionGraphIntoIndependentNodeSubsets メソッドにて、nodes_to_replace から node_subsets に変換。

TfLiteStatus PartitionGraphIntoIndependentNodeSubsets(
    const GraphInfo* info, const TfLiteIntArray* nodes_to_partition,
    std::vector<NodeSubset>* node_subsets) {
  PartitionGraphIntoIndependentNodeSubsetsImpl(info, nodes_to_partition,
                                               node_subsets)
      .Partition();
  return kTfLiteOk;
}

PartitionGraphIntoIndependentNodeSubsetsImpl クラスを生成し、Partitionメソッドを実行。

  // Actually partition the graph.
  void Partition() {
    // Initialize here to make Partition() re-entrant.
    node_subsets_->clear();
    tensor_epochs_.clear();
    tensor_epochs_.resize(info_->num_tensors(), kEpochAlwaysReady);
    node_epochs_.clear();
    node_epochs_.resize(info_->num_nodes(), kEpochNotReady);
    // Set computed tensors to be kEpochNotReady (initializer set everything to
    // AlwaysReady).
    for (int node_index = 0; node_index < info_->num_nodes(); node_index++) {
      const TfLiteNode& node = info_->node(node_index);
      for (int output_tensor_index : TfLiteIntArrayView(node.outputs)) {
        tensor_epochs_[output_tensor_index] = kEpochNotReady;
      }
    }

初期化。。

    // Do a graph traversal where each iteration in the loop is an epoch
    // that corresponds to a node sub set that only contains nodes that are of
    // the same node_type_.
    while (true) {
      BuildNodeSubset();
      if (node_subsets_->back().nodes.empty()) {
        node_subsets_->pop_back();
        break;
      }
    }

    // Mark model outputs as node sub set outputs. All the rest have already
    // been identified.
    for (int output_index : info_->outputs()) {
      int output_epoch = tensor_epochs_[output_index];
      NodeSubset& output_subset = (*node_subsets_)[output_epoch];
      output_subset.output_tensors.push_back(output_index);
    }
    // Make sure every node sub set's inputs and outputs are unique. Since the
    // list of inputs and outputs is generated in a way that produces
    // duplicates.
    for (NodeSubset& node_subset : *node_subsets_) {
      // Sort and uniquefy using standard library algorithms.
      auto uniquefy = [](std::vector<int>* items) {
        std::sort(items->begin(), items->end());
        auto last = std::unique(items->begin(), items->end());
        items->erase(last, items->end());
      };
      uniquefy(&node_subset.input_tensors);
      uniquefy(&node_subset.output_tensors);
    }
  }

NodeSubset を生成。。。

  // Completely populates the current node_subset by doing graph traversal
  void BuildNodeSubset() {
    node_subsets_->emplace_back(NodeSubset());
    // loop until no more nodes can be updated.
    while (true) {
      bool did_something = false;
      for (int node_index = 0; node_index < info_->num_nodes(); node_index++) {
        if (UpdateNode(node_index)) {
          did_something = true;
        }
      }
      if (!did_something) return;
    }
  }

ここまでで、node_subsets に Partition されたノード (NodeSubsetクラス)がストアされている。

  execution_plan_.clear();

  for (auto& node_subset : node_subsets) {
    // Subsets calimed by the delegate should have a "macro" op created, the
    // other node_subsets (kTfNonPartition) just have their nodes added back to
    // the execution plan.
    switch (node_subset.type) {

変換した node_subsets のタイプ別に処理を変えている。現時点では、kTfNonPartition と kTfPartition のみ。

      case NodeSubset::kTfNonPartition:
        for (auto it = node_subset.nodes.begin(); it != node_subset.nodes.end();
             ++it) {
          execution_plan_.push_back(*it);
        }
        break;

kTfNonPartition の場合は、その中にあるものを execution_plan_ に push_back する

      case NodeSubset::kTfPartition: {
        int node_index;

        TfLiteDelegateParams* params =
            CreateDelegateParams(delegate, node_subset);
        TF_LITE_ENSURE_STATUS(AddNodeWithParameters(
            node_subset.input_tensors, node_subset.output_tensors, nullptr, 0,
            params, ®istration, &node_index));

        // Initialize the output tensors's delegate-related fields.
        for (int tensor_index : node_subset.output_tensors) {
          TfLiteTensor* tensor = &tensors_[tensor_index];
          TF_LITE_ENSURE(context_, tensor->delegate == nullptr ||
                                       tensor->delegate == delegate);
          tensor->delegate = delegate;
        }

        // Associate the node with the delegate.
        TfLiteNode* node = &nodes_and_registration_[node_index].first;
        node->delegate = delegate;
      } break;

kTfPartition の時は、AddNodeWithParametersでノードを生成。

      case NodeSubset::kTfUnexplored:
        return kTfLiteError;
        break;
    }
  }
  return kTfLiteOk;
}

明日に続く。