GrnLayer.cs

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Text;
using MyCaffe.basecode;
using MyCaffe.common;
using MyCaffe.param;
 
namespace MyCaffe.layers.tft
{
    public class GrnLayer<T> : Layer<T>
    {
        Layer<T> m_ipSkipLayer = null;
        Layer<T> m_ipFc1 = null;
        Layer<T> m_ipContext = null;
        Layer<T> m_act = null;
        Layer<T> m_ipFc2 = null;
        Layer<T> m_dropout = null;
        Layer<T> m_gate = null;
        Layer<T> m_layerNorm = null;
        Blob<T> m_blobResidual = null;
        Blob<T> m_blobIp1 = null;
        Blob<T> m_blobContext = null;
        Blob<T> m_blobContextAdd = null;
        Blob<T> m_blobIp2 = null;
        Blob<T> m_blobGate = null;
        Blob<T> m_blobGatePlusResidual = null;
        Blob<T> m_blobBtm = null;
        BlobCollection<T> m_colTop = new BlobCollection<T>();
        BlobCollection<T> m_colBtm = new BlobCollection<T>();
 
        public GrnLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.GRN;
 
            if (m_param.grn_param.input_dim != m_param.grn_param.output_dim)
                m_blobResidual = new Blob<T>(cuda, log);
 
            m_blobIp1 = new Blob<T>(cuda, log);
            m_blobIp1.Name = p.name + ".ip1";
            m_blobIp2 = new Blob<T>(cuda, log);
            m_blobIp2.Name = p.name + ".ip2";
            m_blobGate = new Blob<T>(cuda, log);
            m_blobGate.Name = p.name + ".gate";
            m_blobGatePlusResidual = new Blob<T>(cuda, log);
            m_blobGatePlusResidual.Name = p.name + ".gate_p_res";
            m_blobBtm = new Blob<T>(cuda, log);
            m_blobBtm.Name = p.name + ".btm";
        }
 
        protected override void dispose()
        {
            dispose(ref m_blobResidual);
            dispose(ref m_blobIp1);
            dispose(ref m_blobContext);
            dispose(ref m_blobContextAdd);
            dispose(ref m_blobIp2);
            dispose(ref m_blobGate);
            dispose(ref m_blobGatePlusResidual);
            dispose(ref m_blobBtm);
 
            dispose(ref m_ipSkipLayer);
            dispose(ref m_ipFc1);
            dispose(ref m_ipContext);
            dispose(ref m_act);
            dispose(ref m_ipFc2);
            dispose(ref m_dropout);
            dispose(ref m_gate);
            dispose(ref m_layerNorm);
        }
 
        protected override void setup_internal_blobs(BlobCollection<T> col)
        {
            if (col.Count > 0)
                return;
 
            if (m_blobContext != null)
                col.Add(m_blobContext);
            if (m_blobContextAdd != null)
                col.Add(m_blobContextAdd);
            if (m_blobResidual != null)
                col.Add(m_blobResidual);
            col.Add(m_blobIp1);
            col.Add(m_blobIp2);
            col.Add(m_blobGate);
            col.Add(m_blobGatePlusResidual);
        }
 
        public override int MinBottomBlobs 
        {
            get { return 1; }
        }
 
        public override int MaxBottomBlobs
        {
            get { return 2; }
        }
 
        public override int ExactNumTopBlobs
        {
            get { return 1; }
        }
 
        private void addBtmTop(Blob<T> btm, Blob<T> top)
        {
            m_colBtm.Clear();
            m_colBtm.Add(btm);
            m_colTop.Clear();
            m_colTop.Add(top);
        }
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            //-------------------------------------------------------
            // Input conditioning components (Eq.4 in original paper)
            //-------------------------------------------------------
 
            // For a direct residual connection, the dimension of the input must match the dimension of the output,
            // otherwise we need to project the input for creating the residual connection.
            if (m_param.grn_param.input_dim != m_param.grn_param.output_dim)
            {
                if (m_ipSkipLayer == null)
                {
                    LayerParameter ip = new LayerParameter(LayerParameter.LayerType.INNERPRODUCT, m_param.name + ".skip");
                    ip.inner_product_param.num_output = (uint)m_param.grn_param.output_dim;
                    ip.inner_product_param.axis = m_param.grn_param.axis;
                    ip.inner_product_param.weight_filler = m_param.grn_param.weight_filler;
                    ip.inner_product_param.bias_filler = m_param.grn_param.bias_filler;
                    m_ipSkipLayer = Layer<T>.Create(m_cuda, m_log, convertLayerParam(ip, m_param), null);
 
                    addBtmTop(colBottom[0], m_blobResidual);
                    m_ipSkipLayer.Setup(m_colBtm, m_colTop);
                    blobs.Add(m_ipSkipLayer.blobs);
                }
            }
 
            // Create the linear layer for projecting the primary input (across time if necessary)
            if (m_ipFc1 == null)
            {
                LayerParameter ip1 = new LayerParameter(LayerParameter.LayerType.INNERPRODUCT, m_param.name + ".fc1");
                ip1.inner_product_param.num_output = (uint)m_param.grn_param.hidden_dim;
                ip1.inner_product_param.axis = m_param.grn_param.axis;
                ip1.inner_product_param.weight_filler = m_param.grn_param.weight_filler;
                ip1.inner_product_param.bias_filler = m_param.grn_param.bias_filler;
                m_ipFc1 = Layer<T>.Create(m_cuda, m_log, convertLayerParam(ip1, m_param), null);
 
                addBtmTop(colBottom[0], m_blobIp1);
                m_ipFc1.Setup(m_colBtm, m_colTop);
                blobs.Add(m_ipFc1.blobs);
            }
            Blob<T> blobIp1 = m_blobIp1;
 
            // If a context input exists, project the context as well.
            if (colBottom.Count > 1)
            {
                if (m_ipContext == null)
                {
                    LayerParameter ip = new LayerParameter(LayerParameter.LayerType.INNERPRODUCT, m_param.name + ".context");
                    ip.inner_product_param.num_output = (uint)m_param.grn_param.hidden_dim;
                    ip.inner_product_param.axis = m_param.grn_param.axis;
                    ip.inner_product_param.weight_filler = m_param.grn_param.weight_filler;
                    ip.inner_product_param.bias_term = false;
                    m_ipContext = Layer<T>.Create(m_cuda, m_log, convertLayerParam(ip, m_param), null);
                    m_blobContext = new Blob<T>(m_cuda, m_log);
                    m_blobContext.Name = m_param.name + ".ctx";
                    m_blobContextAdd = new Blob<T>(m_cuda, m_log);
                    m_blobContextAdd.Name = m_param.name + ".ctx_add";
 
                    addBtmTop(colBottom[1], m_blobContext);
                    m_ipContext.Setup(m_colBtm, m_colTop);
                    blobs.Add(m_ipContext.blobs);
 
                    m_cuda.add(m_blobContext.count(), m_blobContext.gpu_data, m_blobIp1.gpu_data, m_blobContext.mutable_gpu_data);
                }
                blobIp1 = m_blobContext;
            }
 
            // non-linear activation function applied to the sum of projections.
            if (m_act == null)
            {
                if (m_param.grn_param.activation == param.tft.GrnParameter.ACTIVATION.RELU)
                {
                    LayerParameter act = new LayerParameter(LayerParameter.LayerType.RELU, m_param.name + ".act");
                    m_act = Layer<T>.Create(m_cuda, m_log, convertLayerParam(act, m_param), null);
                }
                else
                {
                    LayerParameter act = new LayerParameter(LayerParameter.LayerType.ELU, m_param.name + ".act");
                    act.elu_param.engine = EngineParameter.Engine.CAFFE;
                    act.elu_param.alpha = 1.0;
                    m_act = Layer<T>.Create(m_cuda, m_log, convertLayerParam(act, m_param), null);
                }
 
                addBtmTop(blobIp1, blobIp1);
                m_act.Setup(m_colBtm, m_colTop);
            }
 
            //-------------------------------------------------------
            // Further projection components (Eq.3 in original paper)
            //-------------------------------------------------------
 
            // Create the linear layer for projecting top of the activation function
            if (m_ipFc2 == null)
            {
                LayerParameter ip2 = new LayerParameter(LayerParameter.LayerType.INNERPRODUCT, m_param.name + ".fc2");
                ip2.inner_product_param.num_output = (uint)m_param.grn_param.output_dim;
                ip2.inner_product_param.axis = m_param.grn_param.axis;
                ip2.inner_product_param.weight_filler = m_param.grn_param.weight_filler;
                ip2.inner_product_param.bias_filler = m_param.grn_param.bias_filler;
                m_ipFc2 = Layer<T>.Create(m_cuda, m_log, convertLayerParam(ip2, m_param), null);
 
                addBtmTop(blobIp1, m_blobIp2);
                m_ipFc2.Setup(m_colBtm, m_colTop);
                blobs.Add(m_ipFc2.blobs);
            }
 
            //-------------------------------------------------------
            // Output gating components (Eq.2 in original paper)
            //-------------------------------------------------------
 
            if (m_param.grn_param.dropout_ratio > 0)
            {
                if (m_dropout == null)
                {
                    LayerParameter drop = new LayerParameter(LayerParameter.LayerType.DROPOUT, m_param.name + ".drop");
                    drop.dropout_param.dropout_ratio = m_param.grn_param.dropout_ratio;
                    m_dropout = Layer<T>.Create(m_cuda, m_log, convertLayerParam(drop, m_param), null);
 
                    addBtmTop(m_blobIp2, m_blobIp2);
                    m_dropout.Setup(m_colBtm, m_colTop);
                }
            }
 
            if (m_gate == null)
            {
                LayerParameter gate = new LayerParameter(LayerParameter.LayerType.GLU, m_param.name + ".gate");
                gate.glu_param.input_dim = m_param.grn_param.output_dim;
                gate.glu_param.axis = m_param.grn_param.axis;
                gate.glu_param.weight_filler = m_param.grn_param.weight_filler;
                gate.glu_param.bias_filler = m_param.grn_param.bias_filler;
                m_gate = Layer<T>.Create(m_cuda, m_log, convertLayerParam(gate, m_param), null);
 
                addBtmTop(m_blobIp2, m_blobGate);
                m_gate.Setup(m_colBtm, m_colTop);
                blobs.Add(m_gate.blobs);
            }
 
            if (m_layerNorm == null)
            {
                LayerParameter layerNorm = new LayerParameter(LayerParameter.LayerType.LAYERNORM, m_param.name + ".layernorm");
                layerNorm.layer_norm_param.epsilon = 1e-10;
                m_layerNorm = Layer<T>.Create(m_cuda, m_log, convertLayerParam(layerNorm, m_param), null);
 
                addBtmTop(m_blobGate, colTop[0]);
                m_layerNorm.Setup(m_colBtm, m_colTop);
            }
 
            setup_internal_blobs(m_colInternalBlobs);
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            if (m_ipSkipLayer != null)
            {
                addBtmTop(colBottom[0], m_blobResidual);
                m_ipSkipLayer.Reshape(m_colBtm, m_colTop);
            }
 
            addBtmTop(colBottom[0], m_blobIp1);
            m_ipFc1.Reshape(m_colBtm, m_colTop);
            Blob<T> blobIp1 = m_blobIp1;
 
            if (colBottom.Count > 1)
            {
                addBtmTop(colBottom[1], m_blobContext);
                m_ipContext.Reshape(m_colBtm, m_colTop);
                m_blobContextAdd.ReshapeLike(m_blobContext);
                blobIp1 = m_blobContext;
            }
 
            addBtmTop(blobIp1, blobIp1);
            m_act.Reshape(m_colBtm, m_colTop);
 
            addBtmTop(blobIp1, m_blobIp2);
            m_ipFc2.Reshape(m_colBtm, m_colTop);
 
            if (m_dropout != null)
            {
                addBtmTop(m_blobIp2, m_blobIp2);
                m_dropout.Reshape(m_colBtm, m_colTop);
            }
 
            addBtmTop(m_blobIp2, m_blobGate);
            m_gate.Reshape(m_colBtm, m_colTop);
 
            m_blobGatePlusResidual.ReshapeLike(m_blobGate);
 
            addBtmTop(m_blobGate, colTop[0]);
            m_layerNorm.Reshape(m_colBtm, m_colTop);
 
            m_blobBtm.ReshapeLike(colBottom[0]);
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            Blob<T> blobResidual = colBottom[0];
 
            if (m_ipSkipLayer != null)
            {
                addBtmTop(colBottom[0], m_blobResidual);
                m_ipSkipLayer.Forward(m_colBtm, m_colTop);
                blobResidual = m_blobResidual;
            }
 
            addBtmTop(colBottom[0], m_blobIp1);
            m_ipFc1.Forward(m_colBtm, m_colTop);
            Blob<T> blobIp1 = m_blobIp1;
 
            if (colBottom.Count > 1)
            {
                addBtmTop(colBottom[1], m_blobContext);
                m_ipContext.Forward(m_colBtm, m_colTop);
 
                m_cuda.add(m_blobContext.count(), m_blobIp1.gpu_data, m_blobContext.gpu_data, m_blobContextAdd.gpu_data);
                blobIp1 = m_blobContextAdd;
            }
 
            // act
            addBtmTop(blobIp1, blobIp1);
            m_act.Forward(m_colBtm, m_colTop);
 
            // Fc2
            addBtmTop(blobIp1, m_blobIp2);
            m_ipFc2.Forward(m_colBtm, m_colTop);
 
            // dropout
            if (m_dropout != null)
            {
                addBtmTop(m_blobIp2, m_blobIp2);
                m_dropout.Forward(m_colBtm, m_colTop);
            }
 
            // gate
            addBtmTop(m_blobIp2, m_blobGate);
            m_gate.Forward(m_colBtm, m_colTop);
 
            // add residual
            m_cuda.add(m_blobGatePlusResidual.count(), m_blobGate.gpu_data, blobResidual.gpu_data, m_blobGatePlusResidual.mutable_gpu_data);
 
            // layernorm
            addBtmTop(m_blobGatePlusResidual, colTop[0]);
            m_layerNorm.Forward(m_colBtm, m_colTop);
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            // layernorm
            addBtmTop(m_blobGatePlusResidual, colTop[0]);
            m_layerNorm.Backward(m_colTop, rgbPropagateDown, m_colBtm);
 
            // add residual
            if (m_ipSkipLayer != null)
                m_blobResidual.CopyFrom(m_blobGatePlusResidual, true);
            else
                colBottom[0].CopyFrom(m_blobGatePlusResidual, true, false, 0, true);
 
            m_blobGate.CopyFrom(m_blobGatePlusResidual, true, false, 0, true);
 
            // gate
            addBtmTop(m_blobIp2, m_blobGate);
            m_gate.Backward(m_colTop, rgbPropagateDown, m_colBtm);
 
            // dropout
            if (m_dropout != null)
            {
                addBtmTop(m_blobIp2, m_blobIp2);
                m_dropout.Backward(m_colTop, rgbPropagateDown, m_colBtm);
            }
 
            Blob<T> blobIp1 = m_blobIp1;
            if (colBottom.Count > 1)
                blobIp1 = m_blobContextAdd;
 
            // Fc2
            addBtmTop(blobIp1, m_blobIp2);
            m_ipFc2.Backward(m_colTop, rgbPropagateDown, m_colBtm);
 
            // act
            addBtmTop(blobIp1, blobIp1);
            m_act.Backward(m_colTop, rgbPropagateDown, m_colBtm);
 
            if (colBottom.Count > 1)
            {
                m_blobContext.CopyFrom(m_blobContextAdd, true);
                m_blobIp1.CopyFrom(m_blobContextAdd, true);
 
                addBtmTop(colBottom[1], m_blobContext);
                m_ipContext.Backward(m_colTop, rgbPropagateDown, m_colBtm);
            }
 
            m_blobBtm.CopyFrom(colBottom[0]);
            addBtmTop(m_blobBtm, m_blobIp1);
            m_ipFc1.Backward(m_colTop, rgbPropagateDown, m_colBtm);
 
            m_cuda.add(colBottom[0].count(), colBottom[0].gpu_diff, m_blobBtm.gpu_diff, colBottom[0].mutable_gpu_diff);
 
            if (m_ipSkipLayer != null)
            {
                addBtmTop(m_blobBtm, m_blobResidual);
                m_ipSkipLayer.Backward(m_colTop, rgbPropagateDown, m_colBtm);
                m_cuda.add(colBottom[0].count(), colBottom[0].gpu_diff, m_blobBtm.gpu_diff, colBottom[0].mutable_gpu_diff);
            }
        }
    }
}