DropoutLayer.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 DropoutLayer<T> : NeuronLayer<T>
    {
        Blob<T> m_blobRand;
        double m_dfThreshold;
        double m_dfScale;
        uint m_uiThreshold;
 
        long m_hCuda = 0;
        long m_hBottomDesc = 0;
        long m_hStates = 0;
        long m_hReserved = 0;
        long m_hDropoutDesc = 0;
        string m_strBottomSize = null;
        ulong m_ulStates = 0;
        ulong m_ulReserved = 0;
 
        public DropoutLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.DROPOUT;
            m_blobRand = new Blob<T>(cuda, log);
            m_blobRand.Name = m_param.name + " rand";
        }
 
        protected override void dispose()
        {
            if (m_blobRand != null)
            {
                m_blobRand.Dispose();
                m_blobRand = null;
            }
 
            if (m_hDropoutDesc != 0)
            {
                m_cuda.FreeDropoutDesc(m_hDropoutDesc);
                m_hDropoutDesc = 0;
            }
 
            if (m_hStates != 0)
            {
                m_cuda.FreeMemory(m_hStates);
                m_hStates = 0;
            }
 
            if (m_hReserved != 0)
            {
                m_cuda.FreeMemory(m_hReserved);
                m_hReserved = 0;
            }
 
            if (m_hBottomDesc != 0)
            {
                m_cuda.FreeTensorDesc(m_hBottomDesc);
                m_hBottomDesc = 0;
            }
 
            if (m_hCuda != 0)
            {
                m_cuda.FreeCuDNN(m_hCuda);
                m_hCuda = 0;
            }
 
            base.dispose();
        }
 
        protected override void setup_internal_blobs(BlobCollection<T> col)
        {
            col.Add(m_blobRand);
        }
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            base.LayerSetUp(colBottom, colTop);
 
            m_dfThreshold = m_param.dropout_param.dropout_ratio;
            m_log.CHECK(m_dfThreshold > 0.0, "Threshold should be > 0");
            m_log.CHECK(m_dfThreshold < 1.0, "Threshold should be < 1");
            m_dfScale = 1.0 / (1.0 - m_dfThreshold);
            m_uiThreshold = (uint)(uint.MaxValue * m_dfThreshold);
 
            shareLayerBlob(m_blobRand, colBottom[0].shape());
 
            if (!m_param.dropout_param.useCudnn())
                return;
 
            m_hCuda = m_cuda.CreateCuDNN();
            m_hBottomDesc = m_cuda.CreateTensorDesc();
            m_hDropoutDesc = m_cuda.CreateDropoutDesc();
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            base.Reshape(colBottom, colTop);
 
            // Setup the cache for random number generation
            m_blobRand.ReshapeLike(colBottom[0]);
 
            if (!m_param.dropout_param.useCudnn())
                return;
 
            string strBottomSize = colBottom[0].ToSizeString();
            m_log.CHECK(strBottomSize == colTop[0].ToSizeString(), "The bottom[0] and top[0] must have the same size!");
 
            if (strBottomSize != m_strBottomSize)
            {
                ulong ulStates;
                ulong ulReserved;
 
                m_cuda.SetTensorDesc(m_hBottomDesc, colBottom[0].num, colBottom[0].channels, colBottom[0].height, colBottom[0].width);
                m_cuda.GetDropoutInfo(m_hCuda, m_hBottomDesc, out ulStates, out ulReserved);
 
                if (ulStates > m_ulStates)
                {
                    if (m_hStates != 0)
                        m_cuda.FreeMemory(m_hStates);
 
                    m_hStates = m_cuda.AllocMemory((long)ulStates);
                    m_ulStates = ulStates;
                }
 
                if (ulReserved > m_ulReserved)
                {
                    if (m_hReserved != 0)
                        m_cuda.FreeMemory(m_hReserved);
 
                    m_hReserved = m_cuda.AllocMemory((long)ulReserved);
                    m_ulReserved = ulReserved;
                }
 
                long lSeed = m_param.dropout_param.seed;
 
                if (lSeed == 0)
                    lSeed = DateTime.Now.Ticks;
 
                m_cuda.SetDropoutDesc(m_hCuda, m_hDropoutDesc, m_dfThreshold, m_hStates, lSeed);
                m_strBottomSize = strBottomSize;
            }
        }
 
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            if (!m_param.dropout_param.useCudnn())
                forward_caffe(colBottom, colTop);
            else
                forward_cudnn(colBottom, colTop);
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            if (!m_param.dropout_param.useCudnn())
                backward_caffe(colTop, rgbPropagateDown, colBottom);
            else
                backward_cudnn(colTop, rgbPropagateDown, colBottom);
        }
 
        protected void forward_caffe(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            long hBottomData = colBottom[0].gpu_data;
            long hTopData = colTop[0].mutable_gpu_data;
            int nCount = colBottom[0].count();
 
            if (m_phase == Phase.TRAIN && m_param.dropout_param.active)
            {
                long hMask = m_blobRand.mutable_gpu_data;
 
                m_cuda.rng_uniform(nCount, m_tZero, convert(uint.MaxValue), hMask);
                // set thresholds
                m_cuda.dropout_fwd(nCount, hBottomData, hMask, m_uiThreshold, convert(m_dfScale), hTopData);
            }
            else
            {
                m_cuda.copy(nCount, hBottomData, hTopData);
            }
        }
 
        protected void backward_caffe(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            if (!rgbPropagateDown[0])
                return;
 
            long hTopDiff = colTop[0].gpu_diff;
            long hBottomDiff = colBottom[0].mutable_gpu_diff;
 
            if (m_phase == Phase.TRAIN && m_param.dropout_param.active)
            {
                long hMask = m_blobRand.gpu_data;
                int nCount = colBottom[0].count();
 
                m_cuda.dropout_bwd(nCount, hTopDiff, hMask, m_uiThreshold, convert(m_dfScale), hBottomDiff);
            }
            else
            {
                m_cuda.copy(colTop[0].count(), hTopDiff, hBottomDiff);
            }
        }
 
        protected void forward_cudnn(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            long hBottomData = colBottom[0].gpu_data;
            long hTopData = colTop[0].mutable_gpu_data;
 
            if (m_phase == Phase.TRAIN)
                m_cuda.DropoutForward(m_hCuda, m_hDropoutDesc, m_hBottomDesc, hBottomData, m_hBottomDesc, hTopData, m_hReserved);
            else
                m_cuda.copy(colBottom[0].count(), hBottomData, hTopData);
        }
 
        protected void backward_cudnn(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            if (!rgbPropagateDown[0])
                return;
 
            long hTopDiff = colTop[0].gpu_diff;
            long hBottomDiff = colBottom[0].mutable_gpu_diff;
 
            if (m_phase == Phase.TRAIN)
                m_cuda.DropoutBackward(m_hCuda, m_hDropoutDesc, m_hBottomDesc, hTopDiff, m_hBottomDesc, hBottomDiff, m_hReserved);
            else
                m_cuda.copy(colTop[0].count(), hTopDiff, hBottomDiff);
        }
    }
}