ニューラルネットワークの隠れ層を増やしてRastrigin functionの精度を上げよう[c++ Arrayfire]

    constexpr int size_hidden_layer = 1;

    // モデル構築

    std::vector<af::array> data; // 評価値
    data.push_back(af::constant(0.0, size_input, size_batch, dtype_t));
    for (int i_hidden_layer = 0; i_hidden_layer < size_hidden_layer; ++i_hidden_layer)
    {
        data.push_back(af::constant(0.0, size_hidden, size_batch, dtype_t));
        data.push_back(af::constant(0.0, size_hidden, size_batch, dtype_t));
    }
    data.push_back(af::constant(0.0, size_output, size_batch, dtype_t));

    std::vector<af::array> grad; // 誤差(傾き)
    grad.push_back(af::constant(0.0, size_input, size_batch, dtype_t));
    for (int i_hidden_layer = 0; i_hidden_layer < size_hidden_layer; ++i_hidden_layer)
    {
        grad.push_back(af::constant(0.0, size_hidden, size_batch, dtype_t));
        grad.push_back(af::constant(0.0, size_hidden, size_batch, dtype_t));
    }
    grad.push_back(af::constant(0.0, size_output, size_batch, dtype_t));

    std::vector<std::shared_ptr<Layer>> layer; // 処理層
    layer.push_back(std::make_shared<FC_layer>(data.at(0), data.at(1), grad.at(0), grad.at(1)));
    for (int i_hidden_layer = 0; i_hidden_layer < size_hidden_layer * 2; i_hidden_layer += 2)
    {
        layer.push_back(std::make_shared<tanhExp_layer>(data.at(i_hidden_layer + 1), data.at(i_hidden_layer + 2), grad.at(i_hidden_layer + 1), grad.at(i_hidden_layer + 2)));
        layer.push_back(std::make_shared<FC_layer>(data.at(i_hidden_layer + 2), data.at(i_hidden_layer + 3), grad.at(i_hidden_layer + 2), grad.at(i_hidden_layer + 3)));
    }

#include <arrayfire.h>
#undef max
#undef min

#include <iostream>
#include <iomanip>
#include <algorithm>
#include <array>
#include <vector>
#include <memory>
#include <random>

//using var_t = double;
//constexpr auto dtype_t = af::dtype::f64;
//constexpr var_t eps = 1e-7; // floatのときは1e-4くらい、doubleのときは1e-7くらい

using var_t = float;
constexpr auto dtype_t = af::dtype::f32;
constexpr var_t eps = 1e-4; // floatのときは1e-4くらい、doubleのときは1e-7くらい

struct Layer
{
    virtual void init() {}          // 勾配の初期化に用いる
    virtual void forward() = 0;     // 順伝播
    virtual void backward() = 0;    // 逆伝播
    virtual void SGD() {}           // 重みの更新
};

struct FC_layer : public Layer
{
    const af::array& x;
    af::array& y;
    af::array& dx;
    const af::array& dy;

    const int size_batch;

    af::array W;
    af::array B;

    af::array dW;
    af::array dB;

    FC_layer(const af::array& x, af::array& y, af::array& dx, const af::array& dy)
        : x(x)
        , y(y)
        , dx(dx)
        , dy(dy)
        , size_batch(x.dims(1))
    {
        // dimsでその次元方向の要素数を得られる
        // Heの初期値
        W = af::randn(y.dims(0), x.dims(0), 1, dtype_t) * std::sqrt(var_t(2) / var_t(x.dims(0)));
        B = af::randn(y.dims(0), 1, 1, dtype_t) * std::sqrt(var_t(2) / var_t(x.dims(0)));

        dW = W * 0.0;
        dB = B * 0.0;
    }

    virtual void init()
    {
        dW = 0.0;
        dB = 0.0;
    }

    virtual void forward()
    {
        // af::tileは各次元の方向に整数倍コピーする
        y = af::matmul(W, x) + af::tile(B, 1, size_batch);
        y.eval();
    }

    virtual void backward()
    {
        dy.eval();

        // af::matmulNTは2番目の引数の行列を転置して積を求める
        // ミニバッチ分の勾配が全て集計されるので、ミニバッチ数が大きくなるとdW・dBに入る値も大きくなる
        // 今か後か、ミニバッチ数で割ることが必要
        dW += af::matmulNT(dy, x) / (var_t)size_batch;
        dB += af::sum(dy, 1) / (var_t)size_batch;

        dx += af::matmulTN(W, dy);

        dW.eval();
        dB.eval();
    }

    virtual void SGD()
    {
        //constexpr var_t alpha = (var_t)0.01;
        constexpr var_t alpha = (var_t)0.0005;
        W -= alpha * dW;
        B -= alpha * dB;

        W.eval();
        B.eval();
    }
};

struct ReLU_layer : public Layer
{
    const af::array& x;
    af::array& y;
    af::array& dx;
    const af::array& dy;

    ReLU_layer(const af::array& x, af::array& y, af::array& dx, const af::array& dy)
        : x(x)
        , y(y)
        , dx(dx)
        , dy(dy)
    {
    }

    virtual void forward()
    {
        y = af::select((x >= 0.0), x, 0.0); // 要素ごとに、x>=0ならばx,そうでないなら0.0
        y.eval();
    }

    virtual void backward()
    {
        dy.eval();
        dx += af::select((x >= 0.0), dy, 0.0); // 要素ごとに、x>=0ならばdy,そうでないなら0.0
    }
};

struct tanhExp_layer : public Layer
{
    const af::array& x;
    af::array& y;
    af::array& dx;
    const af::array& dy;

    tanhExp_layer(const af::array& x, af::array& y, af::array& dx, const af::array& dy)
        : x(x)
        , y(y)
        , dx(dx)
        , dy(dy)
    {
    }

    inline void forward()
    {
        y = x * af::tanh(af::exp(x));
        y.eval();
    }

    inline void backward()
    {
        dy.eval();
        const af::array x_exp = af::exp(x);
        const af::array tanh_exp = af::tanh(x_exp);
        dx += dy * (tanh_exp - x * x_exp * (tanh_exp * tanh_exp - (var_t)1));
    }
};

int main()
{
    af::info();

    constexpr int size_input = 2;
    constexpr int size_hidden = 300;
    constexpr int size_output = 1;
    constexpr int size_batch = 128;

    constexpr int size_hidden_layer = 6;

    constexpr int size_data = 100000;
    constexpr int size_data_train = (int)(size_data * 0.7);

    // Rastrigin function
    af::array input = af::randu(size_input, size_data, dtype_t) * (var_t)10.24 - (var_t)5.12;
    af::array output = (var_t)(10 * size_input) + af::sum(af::pow2(input) - (var_t)10 * af::cos((var_t)2 * af::Pi * input), 0);


    // モデル構築

    std::vector<af::array> data; // 評価値
    data.push_back(af::constant(0.0, size_input, size_batch, dtype_t));
    for (int i_hidden_layer = 0; i_hidden_layer < size_hidden_layer; ++i_hidden_layer)
    {
        data.push_back(af::constant(0.0, size_hidden, size_batch, dtype_t));
        data.push_back(af::constant(0.0, size_hidden, size_batch, dtype_t));
    }
    data.push_back(af::constant(0.0, size_output, size_batch, dtype_t));

    std::vector<af::array> grad; // 誤差(傾き)
    grad.push_back(af::constant(0.0, size_input, size_batch, dtype_t));
    for (int i_hidden_layer = 0; i_hidden_layer < size_hidden_layer; ++i_hidden_layer)
    {
        grad.push_back(af::constant(0.0, size_hidden, size_batch, dtype_t));
        grad.push_back(af::constant(0.0, size_hidden, size_batch, dtype_t));
    }
    grad.push_back(af::constant(0.0, size_output, size_batch, dtype_t));

    std::vector<std::shared_ptr<Layer>> layer; // 処理層
    layer.push_back(std::make_shared<FC_layer>(data.at(0), data.at(1), grad.at(0), grad.at(1)));
    for (int i_hidden_layer = 0; i_hidden_layer < size_hidden_layer * 2; i_hidden_layer += 2)
    {
        layer.push_back(std::make_shared<tanhExp_layer>(data.at(i_hidden_layer + 1), data.at(i_hidden_layer + 2), grad.at(i_hidden_layer + 1), grad.at(i_hidden_layer + 2)));
        layer.push_back(std::make_shared<FC_layer>(data.at(i_hidden_layer + 2), data.at(i_hidden_layer + 3), grad.at(i_hidden_layer + 2), grad.at(i_hidden_layer + 3)));
    }

    // gradient checking
    {
        // 入力データの設定まで
        auto init_phase = [&]()
        {
            // 誤差(傾き)の初期化
            for (auto& itm : grad) { itm = 0.0; }
            for (auto& itm : layer) { itm->init(); }

            // 入力データの設定
            data.front() = input(af::span, af::seq(0, size_batch - 1));
        };

        // 順伝播と逆伝播
        auto forback = [&]()
        {
            // 順伝播
            for (auto& itm : layer) { itm->forward(); }

            // 誤差の計算
            grad.back() = data.back() - output(af::span, af::seq(0, size_batch - 1)); // y - t

            //逆伝播
            std::for_each(layer.rbegin(), layer.rend(), [](auto& itm) { itm->backward(); });
        };

        var_t grad_test = (var_t)123;
        {
            // 入力データの設定まで
            init_phase();

            // 順伝播と逆伝播
            forback();

            grad_test = af::sum<var_t>(grad.front()(0, 0));
        }

        var_t grad_plus = (var_t)123;
        {
            // 入力データの設定まで
            init_phase();

            // 誤差を加える
            data.front()(0, 0) += eps;

            // 順伝播と逆伝播
            forback();

            grad_plus = af::sum<var_t>(af::pow(grad.back(), 2.0)(0, 0) / (var_t)2); // mse
        }

        var_t grad_minus = (var_t)123;
        {
            // 入力データの設定まで
            init_phase();

            // 誤差を加える
            data.front()(0, 0) -= eps;

            // 順伝播と逆伝播
            forback();

            grad_minus = af::sum<var_t>(af::pow(grad.back(), 2.0)(0, 0) / (var_t)2); // mse
        }

        auto gradapprox = (grad_plus - grad_minus) / (2.0 * eps); // 数値微分
        auto numerator = sqrt(pow(grad_test - gradapprox, 2));
        auto norm_grad = sqrt(grad_test * grad_test);
        auto norm_gradapprox = sqrt(gradapprox * gradapprox);
        auto denominator = norm_grad + norm_gradapprox;
        auto difference = numerator / denominator; // 検証すべき差分

        std::cout
            << "--------" << std::endl
            << std::setw(12) << gradapprox << "" << std::endl
            << std::setw(12) << numerator << " = 0" << std::endl
            << std::setw(12) << norm_grad << std::endl
            << std::setw(12) << norm_gradapprox << "" << std::endl
            << std::setw(12) << denominator << std::endl
            << std::setw(12) << difference << " = 0 " << std::endl;

        std::cout << "-------" << std::endl;
    }

    // 学習
    int epoch = 0;
    constexpr int size_data_in_epoch = size_data_train / size_batch;
    while (true)
    {
        ++epoch;

        // ランダム選出
        // シャッフルされたインデックスを作成する
        af::array idx_data;
        {
            af::array vals_data;
            af::array sort_data = af::randu(size_data_train, 1, dtype_t);
            af::sort(vals_data, idx_data, sort_data, 0);
        }

        for (int step = 0; step < size_data_in_epoch; ++step)
        {
            // 誤差(傾き)の初期化
            for (auto& itm : grad) { itm = 0.0; }
            for (auto& itm : layer) { itm->init(); }

            // 今回のステップの学習対象
            // af::seq は範囲を指定するためのもの
            // コンストラクタによって指定される範囲が特殊なので以下を参照のこと
            // https://arrayfire.org/docs/classaf_1_1seq.htm
            af::array idx_target = idx_data(af::seq((step + 0) * size_batch, (step + 1) * size_batch - 1));

            // 入力値を設定
            data.front() = input(af::span, idx_target);

            // 順伝播
            for (auto& itm : layer) { itm->forward(); }

            // 誤差の計算
            grad.back() = data.back() - output(af::span, idx_target); // y - t

            //逆伝播
            std::for_each(layer.rbegin(), layer.rend(), [](auto& itm) { itm->backward(); });

            // 勾配の更新
            for (auto& itm : layer) { itm->SGD(); }
        }

        // 一定期間ごとにログを取る
        if (epoch % 100 == 1)
        {
            auto diff = af::mean<var_t>(af::abs(grad.back()));
            auto norm = af::norm(grad.back());
            std::cout << "epoch : " << epoch << "\t" << "diff : " << diff << "\t" << "norm : " << norm << std::endl;
        }
        else
        {
            std::cout << "epoch : " << epoch << "\r";
        }
    }

    return 0;
}