C++DNNフレームワークの関数設計に関する考察あるいはポエム

概要

• 自作DNNフレームワークを実装すると既存DNNフレームワークのお気持ちが多少わかるようになる
• 自作DNNフレームワークを構築するにあたり、FunctionをC++でどう実装するか考えていた時の自分用メモを紹介

はじめに

久しぶりのブログ更新になってしまった。 最近は自作DNN (Deep Neural Network) フレームワークの実装が楽しい。この自作DNNフレームワークには、kuuという名前をつけた。中二っぽくていい！

具体的にどのようにDNNフレームワークで学習を行うかは、計算グラフの理論、DNNにおけるBackpropagationのアルゴリズム、フレームワークに含まれる個々のアルゴリズムなどを知っているだけでは自明でない。 自分自身、以前はDNNフレームワークを使いながら特にBackpropagationの具体的な実現方法がブラックボックスとなっていたところがあり、自作DNNフレームワークを実装してみることにした。

ゼロから作るDeep Learning ❸ ―フレームワーク編やそのシリーズの存在は知っていたものの、kuuのドラフト版が出来上がって一通り動くようになってからようやく読んだ。 結果的には、作ってから読むことで本に書いてある考察や著者の悩みポイントにをより深く理解できたと思う。 またシリーズのキャッチコピー (？) である「作る経験はコピーできない。」「作るからこそ、見えるモノ。」という言葉にはちょっと感動してしまった。

自作DNNフレームワークを作ること自体は、自分のための勉強であり、世の中に貢献するような取り組みではないと思う。 しかし作った際の悩みポイントを紹介することで、既存DNNフレームワークの利用者がその設計に感じていた疑問や、気持ち悪さや、愚痴が多少減るのではないかと思ったのでブログへ投稿することにした。

例えば私は「なぜPyTorchもChainerもModule/Linkとfunctional/FunctionNodeを分けるのだろう？Linearとかどっちにも定義しなければいけなくて面倒じゃない？」 という疑問を持っていたのだが、それはkuu設計の初期段階で解消された。 学習対象パラメータを保持しておくためのクラスと、forward/backwardで行う操作とを別で定義しないとどうしてもうまく設計できなかった。また役割が違うのだという事を理解すれば、むしろ分けることの方が自然に思えた。

Chainer のVariableNodeとFunctionNodeを相互に呼び出しあって作り出す計算グラフの設計は美しい。 ただそれはPythonだからきれいに実装できたところが大きいように思う 。ゼロから作るDeep Learning ❸ ―フレームワーク編で紹介されている設計もPythonベースである。 C++でDefine-by-Runを実装してみると、 型や相互参照の制約などでなかなか一筋縄には設計できない。

設計で悩んだ（悩んでいる）点は数多くあるが、この記事ではkuuでFunctionを設計していた際のメモを共有する。

Function

前提

DNNフレームワークの実装において、所与の変数を変形する操作をFunctionと呼ぶことにする。 Functionは、forward-pathで行う操作 (以下、forward) とbackward-pathで行う操作 (以下、backward) とがある。Functionの例として、例えばRelu, Linear, Batch Normalizationなどがある。

Function設計の考察

学習対象パラメータはFunctionとは別の機構 (PyTorchのModuleやChainerのLinkなど) で情報を保持している事とする。 この場合に、必要な変数を仮引数として全て受け取ってしまえば、forwardはstaticメソッドかあるいはクラスに属さない関数でよい。 一方、backwardはforwardで利用した実行時の情報をどこかに保存し、それらを用いて処理を行う必要がある。

Functionのforwardとbackwardの実現方法には、少なくとも下記2つのパターンがある。

• [パターンA] クラスのメソッドとしてforwardとbackwardを定義する方法。例えば、ChainerのFunctionNodeはこのタイプである。

  

  パターンA

• [パターンB] forwardとbackwardをバラバラに定義する方法。PyTorchやTensorFlowの (ライブラリ内で実装されている) Functionはこのタイプである。例えば、Pytorchでは名前空間torch::nn::functional配下でforward関数が一元的に定義され、名前空間torch::autograd::generated配下でforward関数に対応するbackwardクラスが定義されている。

  

  パターンB

パターンA

下記はパターンAでFunctionを実装する例である。

struct Func1{
    Tensor static forward(Tensor　x);
    void backward(Tensor grad);
    Tensor operator() (Tensor　x) {
        return forward(x);
    }
};

struct Func2{
    Tensor static forward(Tensor　x);
    void backward(Tensor grad);
    Tensor operator() (Tensor　x) {
        return forward(x);
    }
};

struct Func3{
    Tensor static forward(Tensor　x);
    void backward(Tensor grad);
    Tensor operator() (Tensor　x) {
        return forward(x);
    }
};

この場合にforwardを呼びだす方法は

auto y = Func1::forward(Func2::forward(Func3::forward(x)));

あるいは

auto y = Func1{}(Func2{}(Func3{}(x)));

となる。

Pros.

• パターンBと比較して、forwardとbackwardの対応関係を把握しやすく、新しいFunctionを実装しやすい。

Cons.

• パターンBと比較して、forwardメソッドを呼ぶ方法では記述量が多く可読性が低くなるし、operator()で呼ぶ方法ではわざわざインスタンスを生成しなければならない。[注1]

[注1] Cons.を解決するためにoperator()をクラスから直接呼び出せるようにstatic化したくなる。同じ要望を持っている人がいて、c++標準化で提案されているようだ。

パターンB

下記はパターンBでFunctionを実装する例である。

namespace forward {
Tensor func1(Tensor x);
Tensor func2(Tensor x);
Tensor func3(Tensor x);
}

namespace backward{
struct Func1Backward(Tensor grad){
    void apply(Tensor grad);
};
struct Func2Backward(Tensor grad){
    void apply(Tensor grad);
};
struct Func3Backward(Tensor grad){
    void apply(Tensor grad);
};
}

この場合にforwardの呼び出しは下記のようになる。

namespace F = forward;
auto y = F::func1(F::func2(F::func3(x)));

Pros.

• FunctionとしてのRelu, Linear, Batch Normalizationなどを関数として利用できるのは自然で可読性も高い (個人感)。

Cons.

• forwardとbackwardの関連を人手で書かないといけないため、backward呼び出し時に間違えないよう注意が必要。

Chainerでの実装例 (パターンA)

• Relu

簡潔に呼び出せない問題を解決するために、Chainerでは下記のように、FunctionNode Classのforwardを呼び出す関数を別途定義している。

class ReLU(function_node.FunctionNode):
    ...
    def forward_cpu(self, inputs):
    ...
    def backward(self, indexes, grad_outputs):
    ...

def relu(x):
    y, = ReLU().apply((x,))
    return y

TensorFlowの実装例 (パターンB)

• Reluのforward
• Reluのbackward

backwardだけでなくforwardもstructで定義され、operator()で呼びだす (つまりインスタンスから呼びだす必要がある) 。

PyTorchでの実装例 (パターンB)

• Reluのforward
• Reluのbackward [注2]

[注2] /autograd/generated/Functions.hはPyTorchにより自動生成されるコードのため、本家リポジトリではない場所を参照する。

PyTorch の場合は、Atenで実装されている関数はderivatives.yaml にforwardとbackwardの紐付けを記載している。このderivatives.yamlとtemplates配下の定義 (特にこれやこれ) を利用して、自動生成スクリプトのなかで autograd/generated/Function.h や autograd/generated/Function.cpp を生成する。

例えばtorch::nn::functional::relu()に対応するbackwardは、autograd/generated/Function.hにstruct ReluBackward0やstruct ReluBackward1として定義される。 そして、torch::nn::functional::relu() (forward) が呼び出されると、下記のようにReluBackward0を計算グラフに登録する。

torch/autograd/generated/VariableType_0.cpp

Tensor relu(const Tensor & self) {
  RECORD_FUNCTION("relu", std::vector<c10::IValue>({self}), Node::peek_at_next_sequence_nr());
  auto& self_ = unpack(self, "self", 0);
  std::shared_ptr<ReluBackward0> grad_fn;
  if (compute_requires_grad( self )) {
    grad_fn = std::shared_ptr<ReluBackward0>(new ReluBackward0(), deleteNode);
    grad_fn->set_next_edges(collect_next_edges( self ));
    grad_fn->self_ = SavedVariable(self, false);
  }
  torch::jit::Node* node = nullptr;
  std::shared_ptr<jit::tracer::TracingState> tracer_state;
  if (jit::tracer::isTracing()) {
    tracer_state = jit::tracer::getTracingState();
    at::Symbol op_name;
    op_name = jit::Symbol::fromQualString("aten::relu");
    node = tracer_state->graph->create(op_name, /*num_outputs=*/0);
    jit::tracer::recordSourceLocation(node);
    jit::tracer::addInputs(node, "self", self);
    tracer_state->graph->insertNode(node);
  
    jit::tracer::setTracingState(nullptr);
  }
  #ifndef NDEBUG
  c10::optional<Storage> self__storage_saved =
    self_.has_storage() ? c10::optional<Storage>(self_.storage()) : c10::nullopt;
  c10::intrusive_ptr<TensorImpl> self__impl_saved;
  if (self_.defined()) self__impl_saved = self_.getIntrusivePtr();
  #endif
  auto tmp = ([&]() {
    at::AutoNonVariableTypeMode non_var_type_mode(true);
    return at::relu(self_);
  })();
  auto result = std::move(tmp);
  #ifndef NDEBUG
  if (self__storage_saved.has_value())
    AT_ASSERT(self__storage_saved.value().is_alias_of(self_.storage()));
  if (self__impl_saved) AT_ASSERT(self__impl_saved == self_.getIntrusivePtr());
  #endif
  if (grad_fn) {
      set_history(flatten_tensor_args( result ), grad_fn);
  }
  if (tracer_state) {
    jit::tracer::setTracingState(std::move(tracer_state));
    jit::tracer::addOutput(node, result);
  }
  return result;
}

kuuでの選択

kuuにおいてパターンA とパターンBのどちらを選ぶかを考えた際、呼び出し方が自然な書き方だと思えたパターンBを選択すことにした。 実は最初はパターンAで実装したのだが、呼び出しを簡潔に書けない点がどうしても気になった。PyTorchのC++APIならkuuよりずっときれいに呼び出せていたのが、そもそもこの考察をしたきっかけである。

ゼロから作るDeep Learning ❸ ―フレームワーク編