GateAddNormLayer.cs

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Text;
using System.Threading;
using MyCaffe.basecode;
using MyCaffe.common;
using MyCaffe.param;
 
namespace MyCaffe.layers.tft
{
    public class GateAddNormLayer<T> : Layer<T>
    {
        int m_nBlocks;
        Layer<T> m_dropout = null;
        Layer<T> m_gate = null;
        Layer<T> m_layerNorm = null;
        Blob<T> m_blobResidual = null;
        Blob<T> m_blobDrop = null;
        Blob<T> m_blobGate = null;
        Blob<T> m_blobGateAddResidual = null;
        BlobCollection<T> m_colTop = new BlobCollection<T>();
        BlobCollection<T> m_colBtm = new BlobCollection<T>();
        List<int> m_rgShape = new List<int>(4);
 
        public GateAddNormLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.GATEADDNORM;
 
            if (m_param.dropout_param != null && m_param.dropout_param.dropout_ratio > 0)
            {
                m_blobDrop = new Blob<T>(cuda, log);
                m_blobDrop.Name = p.name + ".drop";
            }
 
            m_blobResidual = new Blob<T>(cuda, log);
            m_blobResidual.Name = p.name + ".residual";
            m_blobGate = new Blob<T>(cuda, log);
            m_blobGate.Name = p.name + ".gate";
            m_blobGateAddResidual = new Blob<T>(cuda, log);
            m_blobGateAddResidual.Name = p.name + ".gateres";
        }
 
        protected override void dispose()
        {
            dispose(ref m_blobResidual);
            dispose(ref m_blobGate);
            dispose(ref m_blobGateAddResidual);
            dispose(ref m_blobDrop);
 
            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_blobDrop != null)
                col.Add(m_blobDrop);
            col.Add(m_blobGate);
            col.Add(m_blobResidual);
        }
 
        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)
        {
            LayerParameter p;
            Blob<T> blobBtm = colBottom[0];
 
            if (colBottom.Count > 1)
            {
                if (m_param.gateaddnorm_param.residual_channel_offset > 0)
                {
                    int nDiff = colBottom[1].channels - m_param.gateaddnorm_param.residual_channel_offset;
                    if (colBottom[1].channels % nDiff != 0)
                        m_log.FAIL("The number bottom(1).channels must be divisible by the bottom(1).channels - the residual channel offset. For example if bottom(1).channels = 120 and redidual_channel_offset = 90, the difference = 30 which is a factor of both 120 and 90.");
                }
            }
 
            if (m_param.dropout_param != null && m_param.dropout_param.dropout_ratio > 0)
            {
                if (m_dropout == null)
                {
                    p = new LayerParameter(LayerParameter.LayerType.DROPOUT, m_param.name + ".drop");
                    p.dropout_param.Copy(m_param.dropout_param);
                    m_dropout = Layer<T>.Create(m_cuda, m_log, convertLayerParam(p, m_param), null);
 
                    addBtmTop(colBottom[0], m_blobDrop);
                    m_dropout.Setup(m_colBtm, m_colTop);
                }
                blobBtm = m_blobDrop;
            }
 
            if (m_gate == null)
            {
                p = new LayerParameter(LayerParameter.LayerType.GLU, m_param.name + ".glu");
                p.glu_param.Copy(m_param.glu_param);
                m_gate = Layer<T>.Create(m_cuda, m_log, convertLayerParam(p, m_param), null);
 
                addBtmTop(blobBtm, m_blobGate);
                m_gate.Setup(m_colBtm, m_colTop);
                blobs.Add(m_gate.blobs);
            }
            m_blobGateAddResidual.ReshapeLike(m_blobGate);
 
            if (m_layerNorm == null)
            {
                p = new LayerParameter(LayerParameter.LayerType.LAYERNORM, m_param.name + ".layernorm");
                p.layer_norm_param.Copy(m_param.layer_norm_param);
                m_layerNorm = Layer<T>.Create(m_cuda, m_log, convertLayerParam(p, 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)
        {
            Blob<T> blobBtm = colBottom[0];
 
            if (colBottom.Count > 1)
            {
                if (m_param.gateaddnorm_param.residual_channel_offset > 0)
                {
                    int nDiff = colBottom[1].channels - m_param.gateaddnorm_param.residual_channel_offset;
                    m_log.CHECK_EQ(colBottom[1].channels % nDiff, 0, "The bottom(1).channels must be divisible by the bottom(1).channels - residual_channel_offset!");
                    m_nBlocks = colBottom[1].channels / nDiff;
 
                    int nQTimeSteps = nDiff;
                    m_rgShape.Clear();
                    m_rgShape.Add(colBottom[0].num);
                    m_rgShape.Add(nQTimeSteps);
                    m_rgShape.Add(colBottom[0].count(2));
                    m_blobResidual.Reshape(m_rgShape);
                }
                else
                {
                    m_blobResidual.ReshapeLike(colBottom[1]);
                }
            }
 
            if (m_dropout != null)
            {
                addBtmTop(colBottom[0], m_blobDrop);
                m_dropout.Reshape(m_colBtm, m_colTop);
                blobBtm = m_blobDrop;
            }
 
            addBtmTop(blobBtm, m_blobGate);
            m_gate.Reshape(m_colBtm, m_colTop);
            m_blobGateAddResidual.ReshapeLike(m_blobGate);
 
            addBtmTop(m_blobGate, colTop[0]);
            m_layerNorm.Reshape(m_colBtm, m_colTop);
        }
 
        private void copy_to_fwd(BlobCollection<T> colBtm, int nIdx, Blob<T> bTop)
        {
            if (nIdx >= colBtm.Count)
                return;
 
            Blob<T> bBtm = colBtm[nIdx];
 
            if (m_param.gateaddnorm_param.residual_channel_offset > 0)
            {
                // Copy just the future items to the top, so if future = 30,
                // with input shape is btm(256,120,64) just the last (256,30,64) are copied to top 
                int nOuterNum = bBtm.num;
                int nChannels = m_nBlocks;
                int nInnerNum = (bBtm.channels / m_nBlocks) * bBtm.count(2);
                m_cuda.channel_copy(bTop.count(), nOuterNum, nChannels, m_nBlocks, nInnerNum, m_nBlocks-1, bBtm.gpu_data, bTop.mutable_gpu_data, DIR.FWD);
            }
            else
            {
                bTop.CopyFrom(bBtm);
            }
        }
 
        private void copy_to_bwd(BlobCollection<T> colBtm, int nIdx, Blob<T> bTop)
        {
            if (nIdx >= colBtm.Count)
                return;
 
            Blob<T> bBtm = colBtm[nIdx];
 
            if (m_param.gateaddnorm_param.residual_channel_offset > 0)
            {
                // Copy just the future items to the top, so if future = 30,
                // with input shape is btm(256,120,64) just the last (256,30,64) are copied to top 
                int nOuterNum = bBtm.num;
                int nChannels = m_nBlocks;
                int nInnerNum = (bBtm.channels / m_nBlocks) * bBtm.count(2);
                m_cuda.channel_copy(bTop.count(), nOuterNum, nChannels, m_nBlocks, nInnerNum, m_nBlocks - 1, bBtm.mutable_gpu_diff, bTop.gpu_diff, DIR.BWD);
            }
            else
            {
                bBtm.CopyFrom(bTop, true);
            }
        }
 
        private void add_to_bwd(BlobCollection<T> colBtm, int nIdx, Blob<T> bTop)
        {
            if (nIdx >= colBtm.Count)
                return;
 
            Blob<T> bBtm = colBtm[nIdx];
 
            if (m_param.gateaddnorm_param.residual_channel_offset > 0)
            {
                // Copy just the future items to the top, so if future = 30,
                // with input shape is btm(256,120,64) just the last (256,30,64) are copied to top 
                int nOuterNum = bBtm.num;
                int nChannels = m_nBlocks;
                int nInnerNum = (bBtm.channels / m_nBlocks) * bBtm.count(2);
                m_cuda.channel_add(bTop.count(), nOuterNum, nChannels, m_nBlocks, nInnerNum, m_nBlocks - 1, bBtm.mutable_gpu_diff, bTop.gpu_diff, DIR.BWD);
            }
            else
            {
                m_cuda.add(bTop.count(), bTop.gpu_diff, bBtm.gpu_diff, bBtm.mutable_gpu_diff);
            }
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            Blob<T> blobBtm = colBottom[0];
            copy_to_fwd(colBottom, 1, m_blobResidual);
 
            if (m_dropout != null)
            {
                addBtmTop(colBottom[0], m_blobDrop);
                m_dropout.Forward(m_colBtm, m_colTop);
                blobBtm = m_blobDrop;
            }
 
            addBtmTop(blobBtm, m_blobGate);
            m_gate.Forward(m_colBtm, m_colTop);
 
            if (colBottom.Count > 1)
                m_cuda.add(m_blobGateAddResidual.count(), m_blobGate.gpu_data, m_blobResidual.gpu_data, m_blobGateAddResidual.mutable_gpu_data);
            else
                m_blobGateAddResidual.CopyFrom(m_blobGate);
 
            addBtmTop(m_blobGateAddResidual, colTop[0]);
            m_layerNorm.Forward(m_colBtm, m_colTop);
 
            colTop[0].ReshapeLike(m_blobGate);
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            addBtmTop(m_blobGateAddResidual, colTop[0]);
            m_layerNorm.Backward(m_colTop, rgbPropagateDown, m_colBtm);
 
            // Copy grad to the residual if it exists.
            copy_to_bwd(colBottom, 1, m_blobGateAddResidual);
            m_blobGate.CopyFrom(m_blobGateAddResidual, true);
            if (colBottom.Count > 1)
                m_blobResidual.CopyFrom(m_blobGateAddResidual, true);
 
            addBtmTop(colBottom[0], m_blobGate);
            m_gate.Backward(m_colTop, rgbPropagateDown, m_colBtm);
 
            if (m_dropout != null)
            {
                addBtmTop(m_blobDrop, colBottom[0]);
                m_dropout.Backward(m_colTop, rgbPropagateDown, m_colBtm);
                colBottom[0].CopyFrom(m_blobDrop, true);
            }
        }
    }
}