ReductionLayer.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using MyCaffe.basecode;
using MyCaffe.common;
using MyCaffe.param;
 
namespace MyCaffe.layers
{
    public class ReductionLayer<T> : Layer<T>
    {
        ReductionParameter.ReductionOp m_op;
        double m_dfCoeff;
        int m_nAxis;
        int m_nNum;
        int m_nDim;
        Blob<T> m_blobSumMultiplier;
 
        public ReductionLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.REDUCTION;
            m_blobSumMultiplier = new Blob<T>(cuda, log);
            m_blobSumMultiplier.Name = m_param.name + " summult";
        }
 
        protected override void dispose()
        {
            base.dispose();
 
            if (m_blobSumMultiplier != null)
            {
                m_blobSumMultiplier.Dispose();
                m_blobSumMultiplier = null;
            }
        }
 
        protected override void setup_internal_blobs(BlobCollection<T> col)
        {
            if (col.Count > 0)
                return;
 
            col.Add(m_blobSumMultiplier);
        }
 
        public override int ExactNumBottomBlobs
        {
            get { return 1; }
        }
 
        public override int ExactNumTopBlobs
        {
            get { return 1; }
        }
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            m_op = m_param.reduction_param.operation;
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            m_nAxis = colBottom[0].CanonicalAxisIndex(m_param.reduction_param.axis);
 
            // In the output, we'll keep all axes up to the reduction axis, but
            //  throw away all after than.
            // Note: currently reducing along non-tail axes is not supported; otherwise,
            //  we'd need to also copy any axes followign an 'end_axis'.
            List<int> rgTopShape = Utility.Clone<int>(colBottom[0].shape(), m_nAxis);
            colTop[0].Reshape(rgTopShape);
            m_nNum = colBottom[0].count(0, m_nAxis);
            m_nDim = colBottom[0].count(m_nAxis);
            m_log.CHECK_EQ(m_nNum, colTop[0].count(), "The 'num' should equal the top[0].count!");
 
            if (m_op == ReductionParameter.ReductionOp.SUM ||
                m_op == ReductionParameter.ReductionOp.MEAN)
            {
                List<int> rgSumMultShape = new List<int>() { m_nDim };
                m_blobSumMultiplier.Reshape(rgSumMultShape);
                m_blobSumMultiplier.SetData(1);
            }
 
            m_dfCoeff = m_param.reduction_param.coeff;
 
            if (m_op == ReductionParameter.ReductionOp.MEAN)
                m_dfCoeff /= m_nDim;
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            long hBottomData = colBottom[0].gpu_data;
            long hMultData = 0;
 
            if (m_blobSumMultiplier.count() > 0)
                hMultData = m_blobSumMultiplier.gpu_data;
 
            T[] rgTop = colTop[0].mutable_cpu_data;
            int nOffset = 0;
            long lPos;
 
            for (int i = 0; i < m_nNum; i++)
            {
                switch (m_op)
                {
                    case ReductionParameter.ReductionOp.SUM:
                    case ReductionParameter.ReductionOp.MEAN:
                        rgTop[i] = m_cuda.dot(m_nDim, hMultData, hBottomData, 0, nOffset);
                        break;
 
                    case ReductionParameter.ReductionOp.ASUM:
                        rgTop[i] = m_cuda.asum(m_nDim, hBottomData, nOffset);
                        break;
 
                    case ReductionParameter.ReductionOp.SUMSQ:
                        rgTop[i] = m_cuda.dot(m_nDim, hBottomData, hBottomData, nOffset, nOffset);
                        break;
 
                    case ReductionParameter.ReductionOp.MIN:
                        rgTop[i] = convert(m_cuda.min(m_nDim, hBottomData, out lPos, nOffset));
                        break;
 
                    case ReductionParameter.ReductionOp.MAX:
                        rgTop[i] = convert(m_cuda.max(m_nDim, hBottomData, out lPos, nOffset));
                        break;
 
                    default:
                        m_log.FAIL("Unknown reduction op: " + m_op.ToString());
                        break;
                }
 
                nOffset += m_nDim;
            }
 
            colTop[0].mutable_cpu_data = rgTop;
 
            if (m_dfCoeff != 1.0)
                m_cuda.scal(m_nNum, m_dfCoeff, colTop[0].mutable_gpu_data);
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            if (!rgbPropagateDown[0])
                return;
 
            // Get bottom_data, if needed.
            long hBottomData = 0;
 
            if (m_op == ReductionParameter.ReductionOp.ASUM ||
                m_op == ReductionParameter.ReductionOp.SUMSQ)
                hBottomData = colBottom[0].gpu_data;
 
            T[] rgTopDiff = colTop[0].update_cpu_diff();
            long hBottomDiff = colBottom[0].mutable_gpu_diff;
            int nOffset = 0;
 
            for (int i = 0; i < m_nNum; i++)
            {
                double dfBottomCoeff = convertD(rgTopDiff[i]) * m_dfCoeff;
 
                switch (m_op)
                {
                    case ReductionParameter.ReductionOp.SUM:
                    case ReductionParameter.ReductionOp.MEAN:
                        m_cuda.set(m_nDim, hBottomDiff, convert(dfBottomCoeff), -1, nOffset);
                        break;
 
                    case ReductionParameter.ReductionOp.ASUM:
                    case ReductionParameter.ReductionOp.MIN:
                    case ReductionParameter.ReductionOp.MAX:
                        m_cuda.sign(m_nDim, hBottomData, hBottomDiff, nOffset, nOffset);
                        m_cuda.scal(m_nDim, dfBottomCoeff, hBottomDiff, nOffset);
                        break;
 
                    case ReductionParameter.ReductionOp.SUMSQ:
                        m_cuda.scale(m_nDim, convert(2 * dfBottomCoeff), hBottomData, hBottomDiff, nOffset, nOffset);
                        break;
 
                    default:
                        m_log.FAIL("Unknown reduction op: " + m_op.ToString());
                        break;
                }
 
                nOffset += m_nDim;
            }
        }
    }
}