LSTMUnitLayer.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using MyCaffe.basecode;
using MyCaffe.common;
using MyCaffe.param;
 
namespace MyCaffe.layers
{
    public class LSTMUnitLayer<T> : Layer<T>
    {
        // The hidden output dimension.
        int m_nHiddenDim;
        Blob<T> m_blobXActs;
 
        public LSTMUnitLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.LSTM_UNIT;
            m_blobXActs = new Blob<T>(cuda, log);
            m_blobXActs.Name = m_param.name + " xacts";
        }
 
        protected override void dispose()
        {
            if (m_blobXActs != null)
            {
                m_blobXActs.Dispose();
                m_blobXActs = null;
            }
 
            base.dispose();
        }
 
        protected override void setup_internal_blobs(BlobCollection<T> col)
        {
            if (col.Count > 0)
                return;
 
            col.Add(m_blobXActs);
        }
 
        public override int ExactNumBottomBlobs
        {
            get { return 3; }
        }
 
        public override int ExactNumTopBlobs
        {
            get { return 2; }
        }
 
        public override bool AllowForceBackward(int nBottomIdx)
        {
            // Can't propagate to sequence contination indicators.
            if (nBottomIdx != 2)
                return true;
 
            return false;
        }
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            int nNumInstances = colBottom[0].shape(1);
 
            for (int i = 0; i < colBottom.Count; i++)
            {
                if (i == 2)
                    m_log.CHECK_EQ(2, colBottom[i].num_axes, "There should be 2 axes at bottom[2]");
                else
                    m_log.CHECK_EQ(3, colBottom[i].num_axes, "There should be 3 axes at bottom[" + i.ToString() + "]");
 
                m_log.CHECK_EQ(1, colBottom[i].shape(0), "The shape(0) at bottom[" + i.ToString() + "] should be 1.");
                m_log.CHECK_EQ(nNumInstances, colBottom[i].shape(1), "The shape(1) at bottom[" + i.ToString() + "] should equal the number of instances (" + nNumInstances.ToString() + ")");                
            }
 
            m_nHiddenDim = colBottom[0].shape(2);
            m_log.CHECK_EQ(4 * m_nHiddenDim, colBottom[1].shape(2), "The bottom[1].shape(2) should equal the 4 * the number of hidden dimensions (4 x " + m_nHiddenDim.ToString() + ")");
            colTop[0].ReshapeLike(colBottom[0]);
            colTop[1].ReshapeLike(colBottom[0]);
            m_blobXActs.ReshapeLike(colBottom[1]);
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            int nCount = colTop[1].count();
            long hC_prev = colBottom[0].gpu_data;
            long hX = colBottom[1].gpu_data;
            long hCont = colBottom[2].gpu_data;
            long hX_acts = m_blobXActs.mutable_gpu_data;
            long hC = colTop[0].mutable_gpu_data;
            long hH = colTop[1].mutable_gpu_data;
            int nXCount = colBottom[1].count();
 
            m_cuda.lstm_unit_fwd(nCount, m_nHiddenDim, nXCount, hX, hX_acts, hC_prev, hCont, hC, hH);
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            m_log.CHECK(!rgbPropagateDown[2], "Cannot backpropagate to sequence indicators.");
 
            if (!rgbPropagateDown[0] && !rgbPropagateDown[1])
                return;
 
            int nCount = colTop[1].count();
            long hC_prev = colBottom[0].gpu_data;
            long hX_acts = m_blobXActs.gpu_data;
            long hCont = colBottom[2].gpu_data;
            long hC = colTop[0].gpu_data;
            long hH = colTop[1].gpu_data;
            long hC_diff = colTop[0].gpu_diff;
            long hH_diff = colTop[1].gpu_diff;
            long hC_prev_diff = colBottom[0].mutable_gpu_diff;
            long hX_acts_diff = m_blobXActs.mutable_gpu_diff;
            int nXCount = colBottom[1].count();
            long hX_diff = colBottom[1].mutable_gpu_diff;
 
            m_cuda.lstm_unit_bwd(nCount, m_nHiddenDim, nXCount, hC_prev, hX_acts, hC, hH, hCont, hC_diff, hH_diff, hC_prev_diff, hX_acts_diff, hX_diff);
        }
    }
}