PReLULayer.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 PReLULayer<T> : NeuronLayer<T>
    {
        bool m_bChannelShared;
        Blob<T> m_blobMultiplier;      // dot multiplier for backward computation of params.
        Blob<T> m_blobBackwardBuff;    // temporary buffer for backward computation.
        Blob<T> m_blobBottomMemory;    // memory for in-place computation.
 
 
        public PReLULayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.PRELU;
            m_blobMultiplier = new Blob<T>(cuda, log);
            m_blobMultiplier.Name = m_param.name + " mult";
            m_blobBackwardBuff = new Blob<T>(cuda, log);
            m_blobBackwardBuff.Name = m_param.name + " backbuf";
            m_blobBottomMemory = new Blob<T>(cuda, log);
            m_blobBottomMemory.Name = m_param.name + " btmmem";
        }
 
        protected override void dispose()
        {
            m_blobBottomMemory.Dispose();
            m_blobBackwardBuff.Dispose();
            m_blobMultiplier.Dispose();
 
            base.dispose();
        }
 
        protected override void setup_internal_blobs(BlobCollection<T> col)
        {
            if (col.Count > 0)
                return;
 
            col.Add(m_blobMultiplier);
            col.Add(m_blobBackwardBuff);
            col.Add(m_blobBottomMemory);
        }
 
        public override bool ReInitializeParameters(WEIGHT_TARGET target)
        {
            base.ReInitializeParameters(target);
 
            if (target == WEIGHT_TARGET.BOTH || target == WEIGHT_TARGET.WEIGHTS)
            {
                FillerParameter fp = m_param.prelu_param.filler;
                if (fp == null)
                    fp = new FillerParameter("constant", 0.25);
 
                Filler<T> filler = Filler<T>.Create(m_cuda, m_log, fp);
                filler.Fill(m_colBlobs[0]);
            }
 
            return true;
        }
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            m_log.CHECK_GE(colBottom[0].num_axes, 2, "Number of axes of bottom must be >= 2");
            PReLUParameter p = m_param.prelu_param;
            int nChannels = colBottom[0].channels;
 
            m_bChannelShared = p.channel_shared;
 
            if (m_colBlobs.Count > 0)
            {
                m_log.WriteLine("Skipping parameter initialization.");
            }
            else
            {
                m_colBlobs = new BlobCollection<T>();
 
                List<int> rgSlopeShape = new List<int>();
                if (!m_bChannelShared)
                    rgSlopeShape.Add(nChannels);
 
                Blob<T> blobSlope = new Blob<T>(m_cuda, m_log);
                blobSlope.Name = m_param.name + " slope";
                blobSlope.type = BLOB_TYPE.INTERNAL;
 
                if (!shareParameter(blobSlope, rgSlopeShape))
                {
                    blobSlope.Reshape(rgSlopeShape);
                    FillerParameter fp = p.filler;
 
                    if (fp == null)
                        fp = new FillerParameter("constant", 0.25);
 
                    Filler<T> filler = Filler<T>.Create(m_cuda, m_log, fp);
                    filler.Fill(blobSlope);
                }
                m_colBlobs.Add(blobSlope);
            }
 
            if (m_bChannelShared)
                m_log.CHECK_EQ(m_colBlobs[0].count(), 1, "Negative slope size is inconsistent with prototxt config.");
            else
                m_log.CHECK_EQ(m_colBlobs[0].count(), nChannels, "Negative slope size is inconsistent with prototxt config.");
 
            // Propagate gradients to the parameters (as directed by backward pass)
            m_rgbParamPropagateDown = new DictionaryMap<bool>(m_colBlobs.Count, true);
 
            List<int> rgShape = new List<int>() { colBottom[0].count(1) };
 
            m_blobMultiplier.Reshape(rgShape);
            m_blobBackwardBuff.Reshape(rgShape);
            m_blobMultiplier.SetData(1.0);
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            if (!reshapeNeeded(colBottom, colTop))
                return;
 
            m_log.CHECK_GE(colBottom[0].num_axes, 2, "Number of axes of bottom blob must be >= 2.");
            colTop[0].ReshapeLike(colBottom[0]);
 
            if (colBottom[0] == colTop[0])
            {
                // For in-place computation.
                m_blobBottomMemory.ReshapeLike(colBottom[0]);
            }
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            long hBottomData = colBottom[0].gpu_data;
            long hTopData = colTop[0].mutable_gpu_data;
            int nCount = colBottom[0].count();
            int nDim = colBottom[0].count(2);
            int nChannels = colBottom[0].channels;
            long hSlopeData = m_colBlobs[0].gpu_data;
            int nDivFactor = m_bChannelShared ? nChannels : 1;
 
            if (colTop[0] == colBottom[0])
                m_cuda.copy(nCount, hBottomData, m_blobBottomMemory.mutable_gpu_data);
 
            m_cuda.prelu_fwd(nCount, nChannels, nDim, hBottomData, hTopData, hSlopeData, nDivFactor);
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            long hBottomData = colBottom[0].gpu_data;
            long hTopDiff = colTop[0].gpu_diff;
            int nCount = colBottom[0].count();
            int nDim = colBottom[0].count(2);
            int nChannels = colBottom[0].channels;
 
            // for In-place computation
            if (colTop[0] == colBottom[0])
                hBottomData = m_blobBottomMemory.gpu_data;
 
            // Propagate to param
            // Since to write bottom diff will affect top diff if top and bottom blobs
            // are identical (in-place computation), we first compute param backward to
            // keep top_diff unchanged.
            if (m_rgbParamPropagateDown[0])
            {
                long hSlopeDiff = m_colBlobs[0].mutable_gpu_diff;
                int nCDim = nChannels * nDim;
 
                m_cuda.prelu_bwd_param(nCDim, colBottom[0].num, colTop[0].offset(1), hTopDiff, hBottomData, m_blobBackwardBuff.mutable_gpu_diff);
 
                if (m_bChannelShared)
                {
                    T dfSum = m_cuda.dot(nCDim, m_blobBackwardBuff.gpu_diff, m_blobMultiplier.gpu_data);
                    m_cuda.add_scalar(m_colBlobs[0].count(), dfSum, hSlopeDiff);
                }
                else
                {
                    m_cuda.gemv(false, nChannels, nDim, m_tOne, m_blobBackwardBuff.gpu_diff, m_blobMultiplier.gpu_data, m_tOne, hSlopeDiff);
                }
            }
 
            // Propagate to bottom
            if (rgbPropagateDown[0])
            {
                long hBottomDiff = colBottom[0].mutable_gpu_diff;
                long hSlopeData = m_colBlobs[0].gpu_data;
                int nDivFactor = m_bChannelShared ? nChannels : 1;
 
                m_cuda.prelu_bwd(nCount, nChannels, nDim, hTopDiff, hBottomData, hBottomDiff, hSlopeData, nDivFactor);
            }
        }
    }
}