SoftmaxCrossEntropyLossLayer.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 SoftmaxCrossEntropyLossLayer<T> : LossLayer<T>
    {
        SoftmaxLayer<T> m_softmaxLayer;
        Blob<T> m_blobSoftmaxOutput;
        Blob<T> m_blobTarget = null;
        Blob<T> m_blobLoss;
        BlobCollection<T> m_colSoftmaxBottomVec = new BlobCollection<T>();
        BlobCollection<T> m_colSoftmaxTopVec = new BlobCollection<T>();
 
        // How to normalize the loss.
        double m_dfNormalizer = 0;
 
 
        public SoftmaxCrossEntropyLossLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.SOFTMAXCROSSENTROPY_LOSS;
            m_blobSoftmaxOutput = new Blob<T>(cuda, log);
            m_blobSoftmaxOutput.Name = m_param.name + " softmax out";
            m_blobLoss = new Blob<T>(cuda, log);
            m_blobLoss.Name = m_param.name + " loss";
 
            LayerParameter param_softmax = p.Clone(false);
            param_softmax.loss_weight.Clear();
            m_softmaxLayer = new SoftmaxLayer<T>(cuda, log, param_softmax);
        }
 
        protected override void dispose()
        {
            m_blobSoftmaxOutput.Dispose();
            m_softmaxLayer.Dispose();
            m_blobLoss.Dispose();
 
            if (m_blobTarget != null)
                m_blobTarget.Dispose();
 
            base.dispose();
        }
 
        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_colSoftmaxBottomVec = new BlobCollection<T>();
            m_colSoftmaxBottomVec.Add(colBottom[0]);
            m_colSoftmaxTopVec = new BlobCollection<T>();
            m_colSoftmaxTopVec.Add(m_blobSoftmaxOutput);
            m_softmaxLayer.Setup(m_colSoftmaxBottomVec, m_colSoftmaxTopVec);
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            base.Reshape(colBottom, colTop);
            int nAxis = colBottom[0].CanonicalAxisIndex(m_param.softmax_param.axis);
            m_nOuterNum = colBottom[0].count(0, nAxis); // batch size
            m_nInnerNum = colBottom[0].count(nAxis); // instance size: |output| == |target|
 
            if (colBottom[0].count() != colBottom[1].count())
            {
                m_log.CHECK_EQ(colBottom[0].count(0, nAxis), colBottom[1].count(0, nAxis), "SOFTMAX_CROSS_ENTROPY_LOSS layer inputs must have the same count, or the target must have 'num' items of indexes.");
 
                // Set the label at the target index = 1.0
                if (m_blobTarget == null)
                {
                    m_blobTarget = new Blob<T>(m_cuda, m_log);
                    m_blobTarget.Name = "full_label";
                }
 
                m_blobTarget.ReshapeLike(colBottom[0]);
            }
 
            m_softmaxLayer.Reshape(m_colSoftmaxBottomVec, m_colSoftmaxTopVec);
            m_blobLoss.ReshapeLike(colBottom[0]);
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            // Set the target data.
            if (m_blobTarget != null)
            {
                int nAxis = colBottom[0].CanonicalAxisIndex(m_param.softmax_param.axis);
                m_log.CHECK_EQ(colBottom[0].count(0, nAxis), colBottom[1].count(0, nAxis), "SOFTMAX_CROSS_ENTROPY_LOSS layer inputs must have the same count, or the target must have 'num' items of indexes.");
                m_blobTarget.SetData(0);
 
                float[] rgfTarget = convertF(colBottom[1].mutable_cpu_data);
                for (int i = 0; i < colBottom[1].num; i++)
                {
                    int nTargetIdx = (int)rgfTarget[i];
                    m_blobTarget.SetData(1.0, m_nInnerNum * i + nTargetIdx);
                }
            }
 
            // The forward pass computes the softmax outputs.
            m_colSoftmaxBottomVec[0] = colBottom[0];
            m_softmaxLayer.Forward(m_colSoftmaxBottomVec, m_colSoftmaxTopVec);
 
            // Compute the loss (negative log likelihood)
            int nCount = colBottom[0].count();
 
            // Stable version of loss computation for input data.
            long hInputData = colBottom[0].gpu_data;
            long hTarget = (m_blobTarget != null) ? m_blobTarget.gpu_data : colBottom[1].gpu_data;
 
            // Since this memory is not used for anything, we use it here to avoid having
            // to allocate the GPU memory to accumulate intermediate results.
            long hLossData = colBottom[0].mutable_gpu_diff;
            long hCountData = (m_blobTarget != null) ? m_blobTarget.mutable_gpu_diff : colBottom[1].mutable_gpu_diff;
 
            m_cuda.sigmoid_cross_entropy_fwd(nCount, hInputData, hTarget, hLossData, false, -1, hCountData);
 
            double dfValidCount = nCount;
            double dfLoss = m_cuda.asum_double(nCount, hLossData);
            m_dfNormalizer = get_normalizer(m_normalization, (int)dfValidCount);
 
            colTop[0].SetData(dfLoss / m_dfNormalizer, 0);
 
            // Return the losses in colTop[1] if it exists.
            if (colTop.Count == 2)
            {
                m_cuda.copy(nCount, hLossData, m_blobLoss.mutable_gpu_data);
                colTop[1].ShareData(m_blobLoss);
            }
 
            // Clear scratch memory to prevent interfering with the backward pass (see #6202)
            colBottom[0].SetDiff(0);
            colBottom[1].SetDiff(0);
 
            if (m_blobTarget != null)
                m_blobTarget.SetDiff(0);
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            if (rgbPropagateDown[1])
                m_log.FAIL(m_type.ToString() + " Layer cannot backpropagate to label inputs.");
 
            if (rgbPropagateDown[0])
            {
                // First, compute the diff.
                int nCount = colBottom[0].count();
                long hSoftmaxOutputData = m_blobSoftmaxOutput.gpu_data;
                long hTarget = (m_blobTarget != null) ? m_blobTarget.gpu_data : colBottom[1].gpu_data;
                long hBottomDiff = colBottom[0].mutable_gpu_diff;
 
                m_cuda.copy(nCount, hSoftmaxOutputData, hBottomDiff);
                m_cuda.axpy(nCount, convert(-1.0), hTarget, hBottomDiff);
 
                // Scale down gradient
                double dfLossWeight = convertD(colTop[0].GetDiff(0)) / m_dfNormalizer;
                m_cuda.scal(nCount, dfLossWeight, hBottomDiff);
            }
        }
    }
}