SoftmaxLayer.cs

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Text;
using MyCaffe.basecode;
using MyCaffe.common;
using MyCaffe.param;
 
namespace MyCaffe.layers
{
    public class SoftmaxLayer<T> : Layer<T>
    {
        int m_nOuterNum;
        int m_nInnerNum;
        int m_nSoftmaxAxis;
        Blob<T> m_blobScale;
        SOFTMAX_ALGORITHM m_algorithm = SOFTMAX_ALGORITHM.DEFAULT;
        SOFTMAX_MODE m_mode = SOFTMAX_MODE.CHANNEL;
        List<int> m_rgScaleDims = null;
 
        long m_hCudnn = 0;
        long m_hBottomDesc = 0;
        long m_hTopDesc = 0;
 
        public SoftmaxLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.SOFTMAX;
            m_blobScale = new Blob<T>(cuda, log);
            m_blobScale.Name = m_param.name + " scale";
 
            setup_internal_blobs(m_colInternalBlobs);
        }
 
        protected override void dispose()
        {
            if (m_hCudnn != 0)
            {
                m_cuda.FreeCuDNN(m_hCudnn);
                m_hCudnn = 0;
            }
 
            if (m_hBottomDesc != 0)
            {
                m_cuda.FreeTensorDesc(m_hBottomDesc);
                m_hBottomDesc = 0;
            }
 
            if (m_hTopDesc != 0)
            {
                m_cuda.FreeTensorDesc(m_hTopDesc);
                m_hTopDesc = 0;
            }
            
            m_blobScale.Dispose();
            base.dispose();
        }
 
        protected override void setup_internal_blobs(BlobCollection<T> col)
        {
            if (col.Count > 0)
                return;
 
            if (!m_param.softmax_param.useCudnn())
            {
                col.Add(m_blobScale);
            }
        }
 
        public override int MinBottomBlobs
        {
            get { return 1; }
        }
 
        public override int MaxBottomBlobs
        {
            get { return 2; }
        }
 
        public override int ExactNumTopBlobs
        {
            get { return 1; }
        }
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            m_algorithm = m_param.softmax_param.algorithm;
            if (m_param.softmax_param.engine == EngineParameter.Engine.CAFFE)
            {
                if (m_algorithm != SOFTMAX_ALGORITHM.ACCURATE)
                    m_log.WriteLine("WARNING: SoftmaxLayer: Caffe mode does not support the ACCURATE algorithm, the default FAST algorithm will be used instead.");
            }
 
            if (!m_param.softmax_param.useCudnn())
                return;
 
            // Initialize cuDNN
            m_hCudnn = m_cuda.CreateCuDNN();
            m_hBottomDesc = m_cuda.CreateTensorDesc();
            m_hTopDesc = m_cuda.CreateTensorDesc();
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            m_nSoftmaxAxis = colBottom[0].CanonicalAxisIndex(m_param.softmax_param.axis);
 
            if (m_phase == Phase.TRAIN && m_param.softmax_param.algorithm_train.HasValue)
                m_algorithm = m_param.softmax_param.algorithm_train.Value;
 
            colTop[0].ReshapeLike(colBottom[0]);
 
            m_nOuterNum = colBottom[0].count(0, m_nSoftmaxAxis);
            m_nInnerNum = colBottom[0].count(m_nSoftmaxAxis + 1);
 
            if (!m_param.softmax_param.useCudnn())
            {
                shareLayerBlob(m_blobScale, colBottom[0].shape());
                m_blobScale.ReshapeLike(colBottom[0]);
 
                m_rgScaleDims = Utility.Clone<int>(colBottom[0].shape());
                m_rgScaleDims[m_nSoftmaxAxis] = 1;
                
                m_blobScale.Reshape(m_rgScaleDims);
 
                return;
            }
 
            int nN = m_nOuterNum;
            int nK = colBottom[0].shape(m_nSoftmaxAxis);
            int nH = m_nInnerNum;
            int nW = 1;
 
            if (nH == 1 && nW == 1)
                m_mode = SOFTMAX_MODE.INSTANCE;
 
            m_cuda.SetTensorDesc(m_hBottomDesc, nN, nK, nH, nW);
            m_cuda.SetTensorDesc(m_hTopDesc, nN, nK, nH, nW);
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            if (!m_param.softmax_param.useCudnn())
                forward_cuda(colBottom, colTop);
            else
                forward_cudnn(colBottom, colTop);
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            if (!m_param.softmax_param.useCudnn())
                backward_cuda(colTop, rgbPropagateDown, colBottom);
            else
                backward_cudnn(colTop, rgbPropagateDown, colBottom);
        }
 
        protected void forward_cuda(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            long hBottomData = colBottom[0].gpu_data;
            long hTopData = colTop[0].mutable_gpu_data;
            long hScaleData = m_blobScale.mutable_gpu_data;
            long hScaleDiff = m_blobScale.mutable_gpu_diff;
            int nCount = colBottom[0].count();
            int nChannels = colTop[0].shape(m_nSoftmaxAxis);
 
            m_cuda.copy(nCount, hBottomData, hTopData);
 
            // We need to subtract the max to avoid numerical issues, compute the exp
            // and then normalize.
            // c = x.max(dim=axis)
            m_cuda.channel_max(m_nOuterNum * m_nInnerNum, m_nOuterNum, nChannels, m_nInnerNum, hTopData, hScaleData);
 
            // Subtract c
            // xm = x - c (along each channel)
            m_cuda.channel_sub(nCount, m_nOuterNum, nChannels, m_nInnerNum, hScaleData, hTopData);
 
            // exponentiate
            // exp_x = exp(xm)
            m_cuda.exp(nCount, hTopData, hTopData);
 
            // Sum across each channel after exp
            // exp_sum = exp_x.sum(dim=axis)
            m_cuda.channel_sum(m_nOuterNum * m_nInnerNum, m_nOuterNum, nChannels, m_nInnerNum, hTopData, hScaleDiff);
 
            if (m_param.softmax_param.algorithm == SOFTMAX_ALGORITHM.LOG)
            {
                // exp_log = exp_sum.log()
                m_cuda.log(m_nOuterNum * m_nInnerNum, hScaleDiff, hScaleDiff);
                // log_z = c + exp_log
                m_cuda.add(m_nOuterNum * m_nInnerNum, hScaleData, hScaleDiff, hScaleData);
                // sm = x - log_z
                m_cuda.copy(nCount, hBottomData, hTopData);
                m_cuda.channel_sub(nCount, m_nOuterNum, nChannels, m_nInnerNum, hScaleData, hTopData);
            }
            else
            {
                // divide            
                m_cuda.channel_div(nCount, m_nOuterNum, nChannels, m_nInnerNum, hScaleDiff, hTopData);
            }
        }
 
 
        protected void backward_cuda(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            long hTopDiff = colTop[0].gpu_diff;
            long hTopData = colTop[0].gpu_data;
            long hBottomDiff = colBottom[0].mutable_gpu_diff;
            long hScaleData = m_blobScale.mutable_gpu_data;
            int nCount = colTop[0].count();
            int nChannels = colTop[0].shape(m_nSoftmaxAxis);
 
            if (m_param.softmax_param.algorithm == SOFTMAX_ALGORITHM.LOG)
            {
                // sumgy = sum channel diff
                m_blobScale.Reshape(m_rgScaleDims);
                m_cuda.channel_sum(nCount, m_nOuterNum, nChannels, m_nInnerNum, hTopDiff, m_blobScale.mutable_gpu_diff);
 
                // expy = exp(y)
                m_cuda.exp(nCount, hTopData, hBottomDiff);
 
                // Fill the expy values across each channel.
                m_blobScale.ReshapeLike(colBottom[0]);
                m_cuda.channel_fillfrom(nCount, m_nOuterNum, 1, nChannels, m_blobScale.gpu_diff, m_blobScale.mutable_gpu_data, DIR.FWD);
 
                // expy * sumgy
                m_cuda.mul(nCount, hBottomDiff, m_blobScale.gpu_data, hBottomDiff);
                m_blobScale.Reshape(m_rgScaleDims);
 
                // grad = gy - (expy * sumgy)
                m_cuda.sub(nCount, hTopDiff, hBottomDiff, hBottomDiff);
            }
            else
            {
                m_cuda.copy(nCount, hTopDiff, hBottomDiff);
 
                // Compute inner1d(top_diff, top_data) and subtract them from the bottom diff.
                m_cuda.channel_dot(m_nOuterNum * m_nInnerNum, m_nOuterNum, nChannels, m_nInnerNum, hTopDiff, hTopData, hScaleData);
                m_cuda.channel_sub(nCount, m_nOuterNum, nChannels, m_nInnerNum, hScaleData, hBottomDiff);
 
                // elementwise multiplication
                m_cuda.mul(nCount, hBottomDiff, hTopData, hBottomDiff);
            }
        }
 
        protected void forward_cudnn(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            long hBottomData = colBottom[0].gpu_data;
            long hTopData = colTop[0].mutable_gpu_data;
 
            m_cuda.SoftmaxForward(m_hCudnn, m_algorithm, m_mode, m_tOne, m_hBottomDesc, hBottomData, m_tZero, m_hTopDesc, hTopData);
        }
 
        protected void backward_cudnn(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            long hTopData = colTop[0].gpu_data;
            long hTopDiff = colTop[0].gpu_diff;
            long hBottomData = colBottom[0].gpu_data;
            long hBottomDiff = colBottom[0].mutable_gpu_diff;
 
            m_cuda.SoftmaxBackward(m_hCudnn, m_algorithm, m_mode, m_tOne, m_hTopDesc, hTopData, m_hTopDesc, hTopDiff, m_tZero, m_hBottomDesc, hBottomDiff);
        }
    }
}