LSTMSimpleLayer.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using MyCaffe.basecode;
using MyCaffe.common;
using MyCaffe.param;
using MyCaffe.fillers;
 
namespace MyCaffe.layers
{
    public class LSTMSimpleLayer<T> : Layer<T>
    {
        int m_nI;   // input dimension.
        int m_nH;   // number of hidden units.
        int m_nT;   // length of sequence.
        int m_nN;   // batch size.
 
        double m_dfClippingThreshold; // threshold for clipped gradient.
        Blob<T> m_blobBiasMultiplier;
 
        Blob<T> m_blobTop;      // Output values.
        Blob<T> m_blobCell;     // Memory cell.
        Blob<T> m_blobPreGate;  // gate values before nonlinearity.
        Blob<T> m_blobGate;     // gate values after nonlinearity.
 
        Blob<T> m_blob_C_0;      // previous cell state value.
        Blob<T> m_blob_H_0;      // previous hidden activation value.
        Blob<T> m_blob_C_T;      // next cell state value.
        Blob<T> m_blob_H_T;      // next hidden activation value.
 
        // Intermediate values.
        Blob<T> m_blob_H_to_Gate;
        Blob<T> m_blob_H_to_H;
 
        public LSTMSimpleLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.LSTM_SIMPLE;
 
            m_blobBiasMultiplier = new Blob<T>(m_cuda, m_log);
            m_blobBiasMultiplier.Name = m_param.name + " biasmult";
            m_blobTop = new Blob<T>(m_cuda, m_log);
            m_blobTop.Name = m_param.name + " top";
            m_blobCell = new Blob<T>(m_cuda, m_log);
            m_blobCell.Name = m_param.name + " cell";
            m_blobPreGate = new Blob<T>(m_cuda, m_log);
            m_blobPreGate.Name = m_param.name + " pregate";
            m_blobGate = new Blob<T>(m_cuda, m_log);
            m_blobGate.Name = m_param.name + " gate";
            m_blob_C_0 = new Blob<T>(m_cuda, m_log);
            m_blob_C_0.Name = m_param.name + " c_0";
            m_blob_H_0 = new Blob<T>(m_cuda, m_log);
            m_blob_H_0.Name = m_param.name + " h_0";
            m_blob_C_T = new Blob<T>(m_cuda, m_log);
            m_blob_C_T.Name = m_param.name + " c_t";
            m_blob_H_T = new Blob<T>(m_cuda, m_log);
            m_blob_H_T.Name = m_param.name + " h_t";
            m_blob_H_to_Gate = new Blob<T>(m_cuda, m_log);
            m_blob_H_to_Gate.Name = m_param.name + "h_to_gate";
            m_blob_H_to_H = new Blob<T>(m_cuda, m_log);
            m_blob_H_to_H.Name = m_param.name + " h_to_h";
        }
 
        protected override void dispose()
        {
            base.dispose();
 
            m_blobBiasMultiplier.Dispose();
            m_blobTop.Dispose();
            m_blobCell.Dispose();
            m_blobPreGate.Dispose();
            m_blobGate.Dispose();
            m_blob_C_0.Dispose();
            m_blob_C_T.Dispose();
            m_blob_H_0.Dispose();
            m_blob_H_T.Dispose();
            m_blob_H_to_Gate.Dispose();
            m_blob_H_to_H.Dispose();
        }
 
        protected override void setup_internal_blobs(BlobCollection<T> col)
        {
            if (col.Count > 0)
                return;
 
            col.Add(m_blobBiasMultiplier);
            col.Add(m_blobTop);
            col.Add(m_blobCell);
            col.Add(m_blobPreGate);
            col.Add(m_blobGate);
            col.Add(m_blob_C_0);
            col.Add(m_blob_H_0);
            col.Add(m_blob_C_T);
            col.Add(m_blob_H_T);
            col.Add(m_blob_H_to_Gate);
            col.Add(m_blob_H_to_H);
        }
 
        public override bool ReInitializeParameters(WEIGHT_TARGET target)
        {
            base.ReInitializeParameters(target);
 
            if (target == WEIGHT_TARGET.BOTH || target == WEIGHT_TARGET.WEIGHTS)
            {
                Filler<T> weight_filler = Filler<T>.Create(m_cuda, m_log, m_param.lstm_simple_param.weight_filler);
                weight_filler.Fill(m_colBlobs[0]);
                weight_filler.Fill(m_colBlobs[1]);
            }
 
            if (target == WEIGHT_TARGET.BOTH || target == WEIGHT_TARGET.BIAS)
            {
                Filler<T> bias_filler = Filler<T>.Create(m_cuda, m_log, m_param.lstm_simple_param.bias_filler);
                bias_filler.Fill(m_colBlobs[2]);
 
                // Initialize the bias for the forget gate to 5.0 as described in the
                // Clockwork RNN paper: 
                // [1] Koutnik, J., Greff, K., Gomez, F., Schmidhuber, J., 'A Clockwork RNN', 2014"
                if (m_param.lstm_simple_param.enable_clockwork_forgetgate_bias)
                {
                    double[] rgBias = convertD(m_colBlobs[2].mutable_cpu_data);
 
                    for (int i = m_nH; i < 2 * m_nH; i++)
                    {
                        rgBias[i] = 5.0;
                    }
 
                    m_colBlobs[2].mutable_cpu_data = convert(rgBias);
                }
            }
 
            return true;
        }
 
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            m_dfClippingThreshold = m_param.lstm_simple_param.clipping_threshold;
            m_nN = (int)m_param.lstm_simple_param.batch_size;              // batch size.
            m_nH = (int)m_param.lstm_simple_param.num_output;              // number of hidden units.
            m_nI = (int)(colBottom[0].count() / colBottom[0].num);         // input dimension.
 
            // Check if we need to set up the weights.
            if (m_colBlobs.Count > 0)
            {
                m_log.WriteLine("Skipping parameter initialization.");
            }
            else
            {
                m_colBlobs = new BlobCollection<T>();
 
                Filler<T> weight_filler = Filler<T>.Create(m_cuda, m_log, m_param.lstm_simple_param.weight_filler);
                Filler<T> bias_filler = Filler<T>.Create(m_cuda, m_log, m_param.lstm_simple_param.bias_filler);
 
                // input-to-hidden weights
                // Initialize the weight.
                List<int> rgShape1 = new List<int>() { 4 * m_nH, m_nI };
                Blob<T> blobWeights_I_H = new Blob<T>(m_cuda, m_log);
                blobWeights_I_H.Name = m_param.name + " weights I to H";
                blobWeights_I_H.type = BLOB_TYPE.WEIGHT;
 
                if (!shareParameter(blobWeights_I_H, rgShape1))
                {
                    blobWeights_I_H.Reshape(rgShape1);
                    weight_filler.Fill(blobWeights_I_H);
                }
                m_colBlobs.Add(blobWeights_I_H);
 
                // hidden-to-hidden weights
                // Initialize the weight.
                List<int> rgShape2 = new List<int>() { 4 * m_nH, m_nH };
                Blob<T> blobWeights_H_H = new Blob<T>(m_cuda, m_log);
                blobWeights_H_H.Name = m_param.name + " weights H to H";
                blobWeights_H_H.type = BLOB_TYPE.WEIGHT;
 
                if (!shareParameter(blobWeights_H_H, rgShape2))
                {
                    blobWeights_H_H.Reshape(rgShape2);
                    weight_filler.Fill(blobWeights_H_H);
                }
                m_colBlobs.Add(blobWeights_H_H);
 
                // If necessary, initialize and fill the bias term.
                List<int> rgShape3 = new List<int>() { 4 * m_nH };
                Blob<T> blobBias = new Blob<T>(m_cuda, m_log);
                blobBias.Name = m_param.name + " bias weights";
                blobBias.type = BLOB_TYPE.WEIGHT;
 
                if (!shareParameter(blobBias, rgShape3))
                {
                    blobBias.Reshape(rgShape3);
                    bias_filler.Fill(blobBias);
                }
                m_colBlobs.Add(blobBias);
 
                // Initialize the bias for the forget gate to 5.0 as described in the
                // Clockwork RNN paper: 
                // [1] Koutnik, J., Greff, K., Gomez, F., Schmidhuber, J., 'A Clockwork RNN', 2014"
                if (m_param.lstm_simple_param.enable_clockwork_forgetgate_bias)
                {
                    double[] rgBias = convertD(blobBias.mutable_cpu_data);
 
                    for (int i=m_nH; i<2*m_nH; i++)
                    {
                        rgBias[i] = 5.0;
                    }
 
                    blobBias.mutable_cpu_data = convert(rgBias);
                }
            }
 
            m_rgbParamPropagateDown = new DictionaryMap<bool>(m_colBlobs.Count, true);
 
            List<int> rgCellShape = new List<int>() { m_nN, m_nH };
            m_blob_C_0.Reshape(rgCellShape);
            m_blob_H_0.Reshape(rgCellShape);
            m_blob_C_T.Reshape(rgCellShape);
            m_blob_H_T.Reshape(rgCellShape);
            m_blob_H_to_H.Reshape(rgCellShape);
 
            List<int> rgGateShape = new List<int>() { m_nN, 4, m_nH };
            m_blob_H_to_Gate.Reshape(rgGateShape);
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            if (m_bNetReshapeRequest)
            {
                m_nN = colBottom[0].num;
                m_bNetReshapeRequest = false;
            }
 
            // Figure out the dimensions.
            m_nT = colBottom[0].num / m_nN;     // length of sequence.
            m_log.CHECK_EQ(colBottom[0].num % m_nN, 0, "The inputs size should be a multiple of the batch size.");
            m_log.CHECK_EQ(colBottom[0].count() / m_nT / m_nN, m_nI, "The input size is incompatible with inner product parameters.");
 
            List<int> rgOriginalTopShape = new List<int>() { m_nT * m_nN, m_nH };
            colTop[0].Reshape(rgOriginalTopShape);
 
            // Gate initialization.
            List<int> rgGateShape = new List<int>() { m_nT, m_nN, 4, m_nH };
            m_blobPreGate.Reshape(rgGateShape);
            m_blobGate.Reshape(rgGateShape);
            m_blob_H_to_Gate.Reshape(rgGateShape);
 
            List<int> rgTopShape = new List<int>() { m_nT, m_nN, m_nH };
            m_blobCell.Reshape(rgTopShape);
            m_blobTop.Reshape(rgTopShape);
            m_blobTop.ShareData(colTop[0]);
            m_blobTop.ShareDiff(colTop[0]);
 
            // Setup the bias multipler.
            List<int> rgMultiplierShape = new List<int>() { m_nN * m_nT };
            m_blobBiasMultiplier.Reshape(rgMultiplierShape);
            m_blobBiasMultiplier.SetData(1.0);
 
            List<int> rgCellShape = new List<int>() { m_nN, m_nH };
            m_blob_C_0.Reshape(rgCellShape);
            m_blob_H_0.Reshape(rgCellShape);
            m_blob_C_T.Reshape(rgCellShape);
            m_blob_H_T.Reshape(rgCellShape);
            m_blob_H_to_H.Reshape(rgCellShape);
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            m_log.CHECK_EQ(colTop[0].gpu_data, m_blobTop.gpu_data, "The top[0].gpu_data should equal the blobTop.gpu_data.");
            long hTopData = m_blobTop.mutable_gpu_data;
            long hBottomData = colBottom[0].gpu_data;
            long hClipData = 0;
 
            if (colBottom.Count > 1)
            {
                hClipData = colBottom[1].gpu_data;
                m_log.CHECK_EQ(colBottom[1].num, colBottom[1].count(), "The bottom[1].num should equal the bottom[1].count.");
            }
 
            long hWeight_i = m_colBlobs[0].gpu_data;
            long hWeight_h = m_colBlobs[1].gpu_data;
            long hBias = m_colBlobs[2].gpu_data;
            long hPreGateData = m_blobPreGate.mutable_gpu_data;
            long hGateData = m_blobGate.mutable_gpu_data;
            long hCellData = m_blobCell.mutable_gpu_data;
            long hHtoGateData = m_blob_H_to_Gate.mutable_gpu_data;
 
            // Initialize previous state.
            if (hClipData != 0)
            {
                m_cuda.copy(m_blob_C_0.count(), m_blob_C_T.gpu_data, m_blob_C_0.mutable_gpu_data);
                m_cuda.copy(m_blob_H_0.count(), m_blob_H_T.gpu_data, m_blob_H_0.mutable_gpu_data);
            }
            else
            {
                m_blob_C_0.SetData(0.0);
                m_blob_H_0.SetData(0.0);
            }
 
            // Compute input to hidden forward propagation.
            m_cuda.gemm(false, true, m_nT * m_nN, 4 * m_nH, m_nI, m_tOne, hBottomData, hWeight_i, m_tZero, hPreGateData);
            m_cuda.gemm(false, false, m_nT * m_nN, 4 * m_nH, 1, m_tOne, m_blobBiasMultiplier.gpu_data, hBias, m_tOne, hPreGateData);
 
            // Compute recurrent forward propagation                
            for (int t = 0; t < m_nT; t++)
            {
                int nTopOffset = m_blobTop.offset(t);
                int nCellOffset = m_blobCell.offset(t);
                int nPreGateOffset = m_blobPreGate.offset(t);
                int nGateOffset = m_blobGate.offset(t);
                int nClipOffset = (hClipData != 0) ? colBottom[1].offset(t) : 0;
                int nHT1Offset;
                long hHT1Data;
                int nCT1Offset;
                long hCT1Data;
 
                if (t == 0)
                {
                    hHT1Data = m_blob_H_0.gpu_data;
                    nHT1Offset = 0;
                    hCT1Data = m_blob_C_0.gpu_data;
                    nCT1Offset = 0;
                }
                else
                {
                    hHT1Data = m_blob_H_T.gpu_data;
                    nHT1Offset = -m_blobTop.offset(1);
                    hCT1Data = m_blob_C_T.gpu_data;
                    nCT1Offset = -m_blobCell.offset(1);
                }
 
                m_cuda.lstm_fwd(t,
                                m_nN,
                                m_nH,
                                m_nI,
                                hWeight_h,
                                hWeight_i,
                                hClipData,
                                nClipOffset,
                                hTopData,       // h_t data
                                nTopOffset,     // h_t offset
                                hCellData,      // c_t data
                                nCellOffset,    // c_t offset
                                hPreGateData,
                                nPreGateOffset,
                                hGateData,
                                nGateOffset,
                                hHT1Data,
                                nHT1Offset,
                                hCT1Data,
                                nCT1Offset,
                                hHtoGateData);
            }
 
            // Preserve cell state and output value for truncated BPTT
            m_cuda.copy(m_nN * m_nH, hCellData, m_blob_C_T.mutable_gpu_data, m_blobCell.offset(m_nT - 1));
            m_cuda.copy(m_nN * m_nH, hTopData, m_blob_H_T.mutable_gpu_data, m_blobTop.offset(m_nT - 1));
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            long hTopData = m_blobTop.gpu_data;
            long hBottomData = colBottom[0].gpu_data;
            long hClipData = 0;
 
            if (colBottom.Count > 1)
            {
                hClipData = colBottom[1].gpu_data;
                m_log.CHECK_EQ(colBottom[1].num, colBottom[1].count(), "The bottom[1].num should equal the bottom[1].count.");
            }
 
            long hWeight_i = m_colBlobs[0].gpu_data;
            long hWeight_h = m_colBlobs[1].gpu_data;
            long hGateData = m_blobGate.gpu_data;
            long hCellData = m_blobCell.gpu_data;
 
            long hTopDiff = m_blobTop.mutable_gpu_diff;
            long hPreGateDiff = m_blobPreGate.mutable_gpu_diff;
            long hGateDiff = m_blobGate.mutable_gpu_diff;
            long hCellDiff = m_blobCell.mutable_gpu_diff;
            long hHtoHData = m_blob_H_to_H.mutable_gpu_data;
 
            m_cuda.copy(m_nN * m_nH, m_blob_C_T.gpu_diff, hCellDiff, 0, m_blobCell.offset(m_nT - 1));
 
            for (int t = m_nT - 1; t >= 0; t--)
            {
                int nTopOffset = m_blobTop.offset(t);
                int nCellOffset = m_blobCell.offset(t);
                int nGateOffset = m_blobGate.offset(t);
                int nPreGateOffset = m_blobPreGate.offset(t);
                int nClipOffset = (hClipData == 0) ? 0 : colBottom[1].offset(t);
                int nCT1Offset;
                long hCT1Data;
                int nDHT1Offset;
                long hDHT1Diff;
                int nDCT1Offset;
                long hDCT1Diff;
 
                if (t == 0)
                {
                    nCT1Offset = 0;
                    hCT1Data = m_blob_C_0.gpu_data;
                    nDHT1Offset = 0;
                    hDHT1Diff = m_blob_H_0.mutable_gpu_diff;
                    nDCT1Offset = 0;
                    hDCT1Diff = m_blob_C_0.mutable_gpu_diff;
                }
                else
                {
                    nCT1Offset = m_blobCell.offset(t - 1);
                    hCT1Data = hCellData;
                    nDHT1Offset = m_blobTop.offset(t - 1);
                    hDHT1Diff = hTopDiff;
                    nDCT1Offset = m_blobCell.offset(t - 1);
                    hDCT1Diff = hCellDiff;
                }
 
                m_cuda.lstm_bwd(t, 
                                m_nN, 
                                m_nH, 
                                m_nI,
                                m_dfClippingThreshold, 
                                hWeight_h, 
                                hClipData, 
                                nClipOffset, 
                                hTopDiff, 
                                nTopOffset, 
                                hCellData, 
                                hCellDiff, 
                                nCellOffset, 
                                hPreGateDiff, 
                                nPreGateOffset, 
                                hGateData, 
                                hGateDiff, 
                                nGateOffset, 
                                hCT1Data, 
                                nCT1Offset, 
                                hDHT1Diff, 
                                nDHT1Offset, 
                                hDCT1Diff, 
                                nDCT1Offset, 
                                hHtoHData);
            }
 
            if (m_rgbParamPropagateDown[0])
            {
                // Gradient w.r.t input-to-hidden weight
                m_cuda.gemm(true, false, 4 * m_nH, m_nI, m_nT * m_nN, m_tOne, hPreGateDiff, hBottomData, m_tOne, m_colBlobs[0].mutable_gpu_diff);
            }
 
            if (m_rgbParamPropagateDown[1])
            {
                // Gradient w.r.t. hidden-to-hidden weight
                m_cuda.gemm(true, false, 4 * m_nH, m_nH, (m_nT - 1) * m_nN, m_tOne, hPreGateDiff, hTopData, m_tOne, m_colBlobs[1].mutable_gpu_diff, m_blobPreGate.offset(1));
 
                // Add gradient from previous time-step.
                m_cuda.gemm(true, false, 4 * m_nH, m_nH, 1, m_tOne, hPreGateDiff, m_blob_H_0.gpu_data, m_tOne, m_colBlobs[1].mutable_gpu_diff);
            }
 
            if (m_rgbParamPropagateDown[2])
            {
                // Gradient w.r.t. bias.
                m_cuda.gemv(true, m_nT * m_nN, 4 * m_nH, m_tOne, hPreGateDiff, m_blobBiasMultiplier.gpu_data, m_tOne, m_colBlobs[2].mutable_gpu_diff);
            }
 
            if (rgbPropagateDown[0])
            {
                // Gradient w.r.t. bottom data.
                m_cuda.gemm(false, false, m_nT * m_nN, m_nI, 4 * m_nH, m_tOne, hPreGateDiff, hWeight_i, m_tZero, colBottom[0].mutable_gpu_diff);
            }
        }
    }
}