GatherLayer.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using MyCaffe.basecode;
using MyCaffe.common;
using MyCaffe.param;
using MyCaffe.param.beta;
 
namespace MyCaffe.layers.beta
{
    public class GatherLayer<T> : Layer<T>
    {
        int m_nAxis = 0;
        int m_nM = 0;
        int m_nN = 0;
        int m_nDim = 0;
        int m_nDimAtAxis = 0;
 
        public GatherLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.GATHER;
        }
 
        public override int ExactNumBottomBlobs
        {
            get { return 2; }
        }
 
        public override int ExactNumTopBlobs
        {
            get { return 1; }
        }
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            m_nAxis = colBottom[0].CanonicalAxisIndex(m_param.gather_param.axis);
 
            if (m_nAxis != 0 && m_nAxis != 1)
                m_log.FAIL("Currently only axis = 0 or axis = 1 are supported.");
 
            m_nDim = colBottom[0].count(m_nAxis + 1);
            m_nDimAtAxis = colBottom[0].shape()[m_nAxis];
            m_nM = colBottom[0].count(0, m_nAxis);
            m_nN = colBottom[1].count();
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            string strErr;
            float[] rgIdxf = convertF(colBottom[1].mutable_cpu_data);
            List<int> rgTopShape = GatherParameter.Reshape(m_nAxis, colBottom[0].shape(), colBottom[1].shape(), rgIdxf, out m_nDim, out m_nDimAtAxis, out m_nM, out m_nN, out strErr);
            if (rgTopShape == null)
                m_log.FAIL(strErr);
 
            colTop[0].Reshape(rgTopShape);
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            int nCount = colBottom[0].count();
            long hBottom = colBottom[0].gpu_data;
            long hIdx = colBottom[1].gpu_data;
            long hTop = colTop[0].mutable_gpu_data;
            int nExpectedCount = (m_nAxis == 0) ? (m_nN * m_nDim) : (m_nN * m_nM);
 
            m_log.CHECK_EQ(colTop[0].count(), nExpectedCount, "The top count should equal " + nExpectedCount.ToString() + "!");
 
            m_cuda.gather_fwd(nCount, hBottom, hTop, m_nAxis, m_nDim, m_nDimAtAxis, m_nM, m_nN, hIdx);
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            int nCount = colBottom[0].count();
            long hBottomDiff = colBottom[0].mutable_gpu_diff;
            long hIdx = colBottom[1].gpu_data;
            long hTopDiff = colTop[0].gpu_diff;
 
            m_cuda.gather_bwd(nCount, hTopDiff, hBottomDiff, m_nAxis, m_nDim, m_nDimAtAxis, m_nM, m_nN, hIdx);
        }
    }
}