たたみ込みニューラルネットをC++11とTBBで実装

動機

PRMLと参考文献読んで、たたみ込みニューラルネットが気になった。世の中の実装を探してみたんだけど、CUDAが必要だったり実装が怪しかったりしたので、勉強がてら自分で書いてみた。

たたみ込みニューラルネットとは

PRML5.5.6に出てくるアレ。1998年の以下の論文がたぶん初出。

Gradient Based Learning Applied to Document Recognition

7層前後の深い階層構造を持っており、最近流行りのdeep learningからすると恐らくご先祖のようなもの。論文中でもdeep learningみたく、「特徴抽出を問題ごとに手探りするのは止めて、生画像を深いネットワークに突っ込んで学習させた方がいいよね」という方向性をはっきり主張している。

The main message of this paper is that better pattern recognition systems can be built by relying more on automatic learning, and less on hand-designed heuristics.
-中略-
Using character recognition as a case study, we show that hand-crafted feature extraction can be advantegeously replaced by carefully designed learning machines that operate directly on pixel images.

Gradient Based Learning Applied to Document Recognition

論文に出てくる手書き数字認識用のLeNet5アーキテクチャはこんな形をしている。

細かい部分を全部省略すると、「5x5のフィルタを複数種類、複数段並べた特徴抽出部と、普通のニューラルネットによる識別部をまるごと全部繋げて、特徴抽出部のカーネル係数と識別部の重みを一緒に学習する」みたいなイメージ。図を眺めると、入力〜C5層までが入力を120次元ベクトルに変換する特徴抽出部、それ以降が3層ニューラルネットによる識別器と見えなくもない。

できたもの

nyanp/tiny-cnn · GitHub
Visual Studio 2012でリリースビルドするか、gccでCNN_USE_TBBを定義してビルドすれば、TBBを使った並列化版が出来る。中身はこんな感じ。

#include <iostream>
#include <boost/timer.hpp>
#include <boost/progress.hpp>

#include "tiny_cnn.h"

using namespace tiny_cnn;

int main(void) {
    // LeNet-5ネットワークの構築
    // RBFネットワークは省略し、6層としている
    typedef network<mse, gradient_descent> CNN;
    CNN nn;
    convolutional_layer<CNN, tanh_activation> C1(32, 32, 5, 1, 6); // input=32x32, window=5, in-channels=1, out-channels=6
    average_pooling_layer<CNN, tanh_activation> S2(28, 28, 6, 2); // input=28x28, input-channels=6, pooling-size=2
    // connection table [Y.Lecun, 1998 Table.1]
    static const bool connection[] = {
        true, false,false,false,true, true, true, false,false,true, true, true, true, false,true, true,
        true, true, false,false,false,true, true, true, false,false,true, true, true, true, false,true,
        true, true, true, false,false,false,true, true, true, false,false,true, false,true, true, true,
        false,true, true, true, false,false,true, true, true, true, false,false,true, false,true, true,
        false,false,true, true, true, false,false,true, true, true, true, false,true, true, false,true,
        false,false,false,true, true, true, false,false,true, true, true, true, false,true, true, true
    };
    convolutional_layer<CNN, tanh_activation> C3(14, 14, 5, 6, 16, connection_table(connection, 6, 16));
    average_pooling_layer<CNN, tanh_activation> S4(10, 10, 16, 2);
    convolutional_layer<CNN, tanh_activation> C5(5, 5, 5, 16, 120);
    fully_connected_layer<CNN, tanh_activation> F6(120, 10);

    nn.add(&C1);
    nn.add(&S2);
    nn.add(&C3);
    nn.add(&S4);
    nn.add(&C5);
    nn.add(&F6);

    // MNISTデータセットをロード
    std::vector<label_t> train_labels, test_labels;
    std::vector<vec_t> train_images, test_images;

    parse_labels("train-labels.idx1-ubyte", &train_labels);
    parse_images("train-images.idx3-ubyte", &train_images);
    parse_labels("t10k-labels.idx1-ubyte", &test_labels);
    parse_images("t10k-images.idx3-ubyte", &test_images);

    boost::progress_display disp(train_images.size());
    boost::timer t;

    // 1epochごとに呼ばれるコールバック
    auto on_enumerate_epoch = [&](){
        std::cout << t.elapsed() << "s elapsed." << std::endl;

        tiny_cnn::result res = nn.test(test_images, test_labels);

        std::cout << nn.learner().alpha << "," << res.num_success << "/" << res.num_total << std::endl;

        nn.learner().alpha *= 0.85;
        nn.learner().alpha = std::max(0.00001, nn.learner().alpha);

        disp.restart(train_images.size());
        t.restart();
    };

    // 1dataごとに呼ばれるコールバック
    auto on_enumerate_data = [&](){ ++disp; };
    
    // training
    nn.init_weight();
    nn.train(train_images, train_labels, 20, on_enumerate_data, on_enumerate_epoch);

    // 結果表示
    nn.test(test_images, test_labels).print_detail(std::cout);

    // 重みをファイルに書き出す
    std::ofstream ofs("LeNet-weights");
    ofs << C1 << S2 << C3 << S4 << C5 << F6;

    //std::ifstream ifs("LeNet-weights");
    //ifs >> C1 >> S2 >> C3 >> S4 >> C5 >> F6;
}

template parameterで評価関数とか活性化関数とかを指定できるようにした。コールバックにlambdaが使えて嬉しい。MNISTのデータセットだと、60000枚の画像の学習が手元PC(4コア2.9GHz)では1周60秒。15分も回せば大体精度99%弱で収束するので、まあまあ速いんじゃないでしょうか。

TBBによる並列化

今回はバッチ学習ではなくstochastic diagonal Levenberg-Marquardtを使っているので、forward-propagation, back-propagationの中の積和計算を細かく並列化している。4コアで2.5倍前後の高速化。TBB有り無しであちこちに#ifdefをバラ撒くのが嫌だったので、以下のような簡単なヘルパを用意した。

typedef tbb::blocked_range<int> blocked_range;

#ifdef CNN_USE_TBB

template<typename Func>
void parallel_for(int begin, int end, Func f) {
    tbb::parallel_for(tbb::blocked_range<int>(begin, end, 100), f); // TBB版
}

#else

template<typename Func>
void parallel_for(int begin, int end, Func f) {
    blocked_range r(begin, end); 
    f(r); // TBB無し版
}

#endif // CNN_USE_TBB

このヘルパ関数を使って、もともとこういう形だったループは、

for (int i = 0; i < this->out_size_; i++) {
    for (int c = 0; c < this->in_size_; c++) 
        l->update(current_delta[i] * prev_out[c], this->Whessian_[i*this->in_size_+c], &this->W_[i*this->in_size_+c]); 

    for (int i = r.begin(); i < r.end(); i++) 
        l->update(current_delta[i], this->bhessian_[i], &this->b_[i]);   
}

少しの変更でTBBあり/無しを切り替えられるようになった。

parallel_for(0,this->out_size_, [&](const blocked_range& r) {
    for (int i = r.begin(); i < r.end(); i++) 
        for (int c = 0; c < this->in_size_; c++) 
            l->update(current_delta[i] * prev_out[c], this->Whessian_[i*this->in_size_+c], &this->W_[i*this->in_size_+c]); 

    for (int i = r.begin(); i < r.end(); i++) 
        l->update(current_delta[i], this->bhessian_[i], &this->b_[i]);   
});

ラムダ+TBB素晴らしい。