NLLLossLayer.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text;
using MyCaffe.basecode;
using MyCaffe.common;
using MyCaffe.param;
 
namespace MyCaffe.layers.gpt
{
    public class NLLLossLayer<T> : LossLayer<T>
    {
        Blob<T> m_blobProb;
        int m_nAxis;
        int? m_nIgnoreLabel = null;
 
        public NLLLossLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.NLL_LOSS;
            m_blobProb = new Blob<T>(m_cuda, m_log);
        }
 
        protected override void dispose()
        {
            dispose(ref m_blobProb);
            base.dispose();
        }
 
        protected override void setup_internal_blobs(BlobCollection<T> col)
        {
            if (col.Count > 0)
                return;
        }
 
        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);
            m_nIgnoreLabel = m_param.loss_param.ignore_label;
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            base.Reshape(colBottom, colTop);
 
            m_blobProb.ReshapeLike(colBottom[0]);
            m_nAxis = colBottom[0].CanonicalAxisIndex(m_param.nll_loss_param.axis);
            m_nOuterNum = colBottom[0].count(0, m_nAxis);
            m_nInnerNum = colBottom[0].count(m_nAxis + 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 nll 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}.");
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            long hProbData = colBottom[0].gpu_data;
            long hLabel = colBottom[1].gpu_data;
            int nDim = colBottom[0].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.nllloss_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 hTopData = colTop[0].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 else,
            // we use to avoid allocating new GPU memory.
            long hCounts = m_blobProb.mutable_gpu_diff;
 
            m_cuda.nllloss_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);
        }
    }
}