SigmoidCrossEntropyLossLayer.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 SigmoidCrossEntropyLossLayer<T> : LossLayer<T>
    {
        SigmoidLayer<T> m_sigmoidLayer;
        Blob<T> m_blobSigmoidOutput;
        Blob<T> m_blobTarget = null;
        Blob<T> m_blobLoss;
        BlobCollection<T> m_colSigmoidBottomVec = new BlobCollection<T>();
        BlobCollection<T> m_colSigmoidTopVec = new BlobCollection<T>();
 
        // The label indicating that an instance should be ignored.
        int? m_nIgnoreLabel = null;
        double m_dfNormalizer = 0;
 
 
        public SigmoidCrossEntropyLossLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.SIGMOIDCROSSENTROPY_LOSS;
            m_blobSigmoidOutput = new Blob<T>(cuda, log);
            m_blobSigmoidOutput.Name = m_param.name + " sigmoid out";
            m_blobLoss = new Blob<T>(cuda, log);
            m_blobLoss.Name = m_param.name + " loss";
 
            LayerParameter param_sigmoid = p.Clone(false);
            param_sigmoid.loss_weight.Clear();
            m_sigmoidLayer = new SigmoidLayer<T>(cuda, log, param_sigmoid);
        }
 
        protected override void dispose()
        {
            m_blobSigmoidOutput.Dispose();
            m_sigmoidLayer.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_colSigmoidBottomVec.Clear();
            m_colSigmoidBottomVec.Add(colBottom[0]);
            m_colSigmoidTopVec.Clear();
            m_colSigmoidTopVec.Add(m_blobSigmoidOutput);
            m_sigmoidLayer.Setup(m_colSigmoidBottomVec, m_colSigmoidTopVec);
 
            m_nIgnoreLabel = m_param.loss_param.ignore_label;
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            base.Reshape(colBottom, colTop);
            m_nOuterNum = colBottom[0].shape(0); // batch size
            m_nInnerNum = colBottom[0].count(1); // instance size: |output| == |target|
 
            if (colBottom[0].count() != colBottom[1].count())
            {
                if (colBottom[1].count() != colBottom[0].num)
                    m_log.FAIL("SIGMOID_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_sigmoidLayer.Reshape(m_colSigmoidBottomVec, m_colSigmoidTopVec);
            m_blobLoss.ReshapeLike(colBottom[0]);
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            // Set the target data.
            if (m_blobTarget != null)
            {
                m_log.CHECK_EQ(colBottom[0].num, colBottom[1].count(), "SIGMOID_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 sigmoid outputs.
            m_colSigmoidBottomVec[0] = colBottom[0];
            m_sigmoidLayer.Forward(m_colSigmoidBottomVec, m_colSigmoidTopVec);
 
            // 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, m_nIgnoreLabel.HasValue, m_nIgnoreLabel.GetValueOrDefault(-1), hCountData);
 
            double dfValidCount = nCount;
            // Only launch another CUDA kernel if we actually need the valid count.
            if (m_normalization == LossParameter.NormalizationMode.VALID && m_nIgnoreLabel.HasValue)
                dfValidCount = m_cuda.asum_double(nCount, hCountData);
 
            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 hSigmoidOutputData = m_blobSigmoidOutput.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, hSigmoidOutputData, hBottomDiff);
                m_cuda.axpy(nCount, convert(-1.0), hTarget, hBottomDiff);
 
                // Zero out gradient for ignored targets
                if (m_nIgnoreLabel.HasValue)
                    m_cuda.sigmoid_cross_entropy_bwd(nCount, m_nIgnoreLabel.Value, hTarget, hBottomDiff);
 
                // Scale down gradient
                double dfLossWeight = convertD(colTop[0].GetDiff(0)) / m_dfNormalizer;
                m_cuda.scal(nCount, dfLossWeight, hBottomDiff);
            }
        }
    }
}