dnn_layer_norm.h

Go to the documentation of this file.

// Copyright 2019 Claes Rolen (www.rolensystems.com)
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.

#pragma once
namespace dnn
{


class layer_norm: public layer
{
private:
    arma::Mat<DNN_Dtype> gamma;
    arma::Mat<DNN_Dtype> beta;
    arma::Mat<DNN_Dtype> x_mean;
    arma::Mat<DNN_Dtype> x_stdinv;
    arma::Mat<DNN_Dtype> x_norm;
    arma::Mat<DNN_Dtype> ee;
    arma::Mat<DNN_Dtype> run_mean;
    arma::Mat<DNN_Dtype> run_var;
    DNN_Dtype run_alpha;
public:
    layer_norm(void):layer() {};

    void init(void)
    {
        layer::init();

        train_par        = true;
        N_cols_right     = left->get_nrof_cols();
        N_rows_right     = left->get_nrof_rows();
        N_channels_right = left->get_nrof_channels();
        N_right          = N_cols_right*N_rows_right*N_channels_right;

        type             = "Batch Norm";
        id               = type;
        gamma.   set_size(N_right, 1); gamma.ones();
        beta.    set_size(N_right, 1); beta.zeros();
        x_stdinv.set_size(N_right, 1);
        ee.      set_size(N_right, 1); ee.fill((DNN_Dtype)1e-5);
        run_mean.set_size(N_right, 1); run_mean.zeros();
        run_var. set_size(N_right, 1); run_var.zeros();
        run_alpha = (DNN_Dtype)0.99;
    }

    void upd_buf_size(arma::uword nmb)
    {
        layer::upd_buf_size(nmb);
        x_mean.set_size(N_right, N_batch); x_mean.zeros();
        x_norm.set_size(N_right, N_batch); x_norm.zeros();
    }

    void set_run_mean(arma::Mat<DNN_Dtype>& r_mean)
    {
        run_mean = r_mean;
    }

    void set_run_var(arma::Mat<DNN_Dtype>& r_var)
    {
        run_var = r_var;
    }

    arma::uword get_nrof_params(void)
    {
        return N_right+N_right;
    }

    void prop(void)
    {
        arma::Mat<DNN_Dtype> X_1 = left->get_Y1();
        Y1 =  gamma%(X_1-run_mean)/arma::sqrt((run_var+ee))+beta;
    }

    void prop_mb(void)
    {
        arma::Mat<DNN_Dtype> x=*(left->get_Y_ptr());

        // Normalize the input.
        const arma::Mat<DNN_Dtype> batch_mean = arma::mean(x,1);
        const arma::Mat<DNN_Dtype> batch_var  = arma::var(x,1,1);
        run_mean = run_mean*run_alpha+batch_mean*((DNN_Dtype)1.0-run_alpha);
        run_var  = run_var*run_alpha+batch_var*((DNN_Dtype)1.0-run_alpha);

        x_mean   = x.each_col()-batch_mean;
        x_stdinv = 1.0/arma::sqrt((batch_var+ee));
        x_norm   = x_mean.each_col()%x_stdinv;

        // Translate output to mean=0, variance=1
        for(arma::uword k=0; k<x_norm.n_cols; k++ )
        {
            Y.col(k) = x_norm.col(k)%gamma + beta;
        }
    }

    void backprop(void)
    {
        arma::Mat<DNN_Dtype> y    =*(right->get_Dleft_ptr());
        arma::Mat<DNN_Dtype> norm = y.each_col()%gamma;
        arma::Mat<DNN_Dtype> var  = -arma::sum(norm % x_mean, 1) % arma::pow(x_stdinv, 3.0)/2;
        Dleft = norm.each_col() % x_stdinv + x_mean.each_col() % (var*((DNN_Dtype)2.0/N_batch));
        Dleft.each_col() += -(arma::sum(norm.each_col() % x_stdinv, 1)- (DNN_Dtype)2.0*(var % arma::mean(x_mean, 1)))/((DNN_Dtype)N_batch);
    }

    void update(void)
    {
        arma::Cube<DNN_Dtype> G(gamma.memptr(),N_right,1,1,false,true);
        arma::Cube<DNN_Dtype> dg(N_right,1,1);
        dg.zeros();
        arma::Mat<DNN_Dtype> db(beta);
        arma::Mat<DNN_Dtype> x  = *(left->get_Y_ptr());
        arma::Mat<DNN_Dtype> dy = *(right->get_Dleft_ptr());

        db = arma::mean(dy,1);
        dg.slice(0).col(0) = arma::sum(x_norm%dy,1);

        opt_alg->apply(G,beta,dg,db);
    }

    void disp(void)
    {
        layer::disp();
        std::cout << "Nr of outputs:        " << get_nrof_outputs() << " ["<< N_rows_right << ","<<N_cols_right<<","<< N_channels_right<<"]"<< std::endl;
    }

}; // End class layer_norm
} // End namespace dnn