SoftmaxLossLayer.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 SoftmaxLossLayer<T> : LossLayer<T>
    {
        Layer<T> m_softmaxLayer;
        Blob<T> m_blobProb;
        BlobCollection<T> m_colSoftmaxBottom;
        BlobCollection<T> m_colSoftmaxTop;
        int? m_nIgnoreLabel = null;
        int m_nSoftmaxAxis;
 
        public SoftmaxLossLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.SOFTMAXWITH_LOSS;
            m_blobProb = new Blob<T>(cuda, log);
            m_blobProb.Name = m_param.name + " prob";
        }
 
        protected override void dispose()
        {
            m_blobProb.Dispose();
 
            if (m_softmaxLayer != null)
                m_softmaxLayer.Dispose();
            
            base.dispose();
        }
 
        protected override void setup_internal_blobs(BlobCollection<T> col)
        {
            if (col.Count > 0)
                return;
 
            col.Add(m_blobProb);
        }
 
        public override int ExactNumTopBlobs
        {
            get { return -1; }
        }
 
        public override int MinTopBlobs
        {
            get { return 1; }
        }
 
        public override int MaxTopBlobs
        {
            get { return 2; }
        }
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            base.LayerSetUp(colBottom, colTop);
 
            LayerParameter param_softmax = m_param.Clone(false);
            param_softmax.SetType(LayerParameter.LayerType.SOFTMAX);
            param_softmax.softmax_param = m_param.softmax_param.Clone() as SoftmaxParameter;
            param_softmax.loss_weight.Clear();
 
            m_softmaxLayer = new SoftmaxLayer<T>(m_cuda, m_log, param_softmax);
            m_colSoftmaxBottom = new BlobCollection<T>();
            m_colSoftmaxTop = new BlobCollection<T>();
 
            m_colSoftmaxBottom.Add(colBottom[0]);
            m_colSoftmaxTop.Add(m_blobProb);
            m_softmaxLayer.Setup(m_colSoftmaxBottom, m_colSoftmaxTop);
 
            m_nIgnoreLabel = m_param.loss_param.ignore_label;
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            base.Reshape(colBottom, colTop);
 
            m_softmaxLayer.Reshape(m_colSoftmaxBottom, m_colSoftmaxTop);
            m_nSoftmaxAxis = colBottom[0].CanonicalAxisIndex(m_param.softmax_param.axis);
            m_nOuterNum = colBottom[0].count(0, m_nSoftmaxAxis);
            m_nInnerNum = colBottom[0].count(m_nSoftmaxAxis + 1);
 
            if (!m_bIgnoreLabels)
            {
                m_log.CHECK_EQ(m_nOuterNum * m_nInnerNum, colBottom[1].count(), "Number of labels must match number of predictions; e.g., if softmax axis == 1 and prediction shape is (N, C, H, W), label count (number of labels) must be N*H*W, with integer values in {0, 1, ..., C-1}.");
 
                if (colTop.Count >= 2)
                {
                    // softmax output
                    colTop[1].ReshapeLike(colBottom[0]);
                }
            }
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            m_softmaxLayer.Forward(m_colSoftmaxBottom, m_colSoftmaxTop);
 
            long hProbData = m_blobProb.gpu_data;
            long hLabel = colBottom[1].gpu_data;
            int nDim = m_blobProb.count() / m_nOuterNum;
            int nCount = m_nOuterNum * m_nInnerNum;
 
            // Since this memory is not used for anything, we use it here to avoid having
            // to allocate new GPU memory to accumulate intermediate results.
            long hLossData = colBottom[0].mutable_gpu_diff;
 
            // Similarly, this memory is never used elsewhere, and thus we can use it
            // to avoid having to allocate additional GPU memory.
            long hCounts = m_blobProb.mutable_gpu_diff;
 
            m_cuda.softmaxloss_fwd(nCount, hProbData, hLabel, hLossData, m_nOuterNum, nDim, m_nInnerNum, hCounts, m_nIgnoreLabel);
            T fLoss = m_cuda.asum(nCount, hLossData);
            double dfValidCount = -1;
 
            // Only launch another cuda kernel if we actually need the count of valid
            // outputs.
            if (m_normalization == LossParameter.NormalizationMode.VALID && m_nIgnoreLabel.HasValue)
                dfValidCount = convertD(m_cuda.asum(nCount, hCounts));
 
            double dfLoss = convertD(fLoss);
            double dfNormalizer = get_normalizer(m_normalization, (int)dfValidCount);
            double dfFinalLoss = dfLoss / dfNormalizer;
 
            colTop[0].SetData(dfFinalLoss, 0);
 
            if (colTop.Count == 2)
                colTop[1].ShareData(m_blobProb);
 
            // Clear scratch memory to prevent with interfering with backward pass (see #602)
            colBottom[0].SetDiff(0);
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            if (!rgbPropagateDown[0])
                return;
 
            long hBottomDiff = colBottom[0].mutable_gpu_diff;
            long hProbData = m_blobProb.gpu_data;
            long hTopData = colTop[0].gpu_data;
 
            m_cuda.copy(m_blobProb.count(), hProbData, hBottomDiff);
 
            long hLabel = colBottom[1].gpu_data;
            int nDim = m_blobProb.count() / m_nOuterNum;
            int nCount = m_nOuterNum * m_nInnerNum;
 
            // Since this memory is not used for anything else,
            // we use to avoid allocating new GPU memory.
            long hCounts = m_blobProb.mutable_gpu_diff;
 
            m_cuda.softmaxloss_bwd(nCount, hTopData, hLabel, hBottomDiff, m_nOuterNum, nDim, m_nInnerNum, hCounts, m_nIgnoreLabel);
 
            double dfValidCount = -1;
 
            // Only launch another cuda kernel if we acutally need the count of valid
            // outputs.
            if (m_normalization == LossParameter.NormalizationMode.VALID && m_nIgnoreLabel.HasValue)
                dfValidCount = convertD(m_cuda.asum(nCount, hCounts));
 
            double dfTopDiff = convertD(colTop[0].GetDiff(0));
            double dfNormalizer = get_normalizer(m_normalization, (int)dfValidCount);
            double dfLossWeight = dfTopDiff / dfNormalizer;
 
            m_cuda.scal(m_blobProb.count(), convert(dfLossWeight), hBottomDiff);
        }
    }
}