InfogainLossLayer.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 InfogainLossLayer<T> : LossLayer<T>
    {
        // The internal SoftmaxLayer used to map predictions to a distribution.
        SoftmaxLayer<T> m_softmaxLayer;
        // Stores the output probability preductions from the SoftmaxLayer.
        Blob<T> m_blobProb = null;
        // Bottom vector holder used in the call to the underlying SoftmaxLayer::Forward
        BlobCollection<T> m_colSoftmaxBottomVec = new BlobCollection<T>();
        // Top vector holder used in the call to the underlying SoftmaxLayer::Forward
        BlobCollection<T> m_colSoftmaxTopVec = new BlobCollection<T>();
        // The infogain blob.
        Blob<T> m_blobInfoGain = null;
        // Cache holding the rows sums of H.
        Blob<T> m_blobSumRowsOfH = null;
        // The optional lable indicating that an instance should be ignored.
        int? m_nIgnoreLabel = null;
        int m_nInfogainAxis = 0;
        int m_nNumLabels = 0;
 
        public InfogainLossLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.INFOGAIN_LOSS;
 
            m_blobProb = new Blob<T>(m_cuda, m_log);
            m_blobProb.Name = m_param.name + " prob";
            m_blobInfoGain = new Blob<T>(m_cuda, m_log);
            m_blobInfoGain.Name = m_param.name + " infogain";
            m_blobSumRowsOfH = new Blob<T>(m_cuda, m_log);
            m_blobSumRowsOfH.Name = m_param.name + " sum rows of H";
        }
 
        protected override void dispose()
        {
            base.dispose();
 
            if (m_blobInfoGain != null)
            {
                m_blobInfoGain.Dispose();
                m_blobInfoGain = null;
            }
 
            if (m_blobProb != null)
            {
                m_blobProb.Dispose();
                m_blobProb = null;
            }
 
            if (m_blobSumRowsOfH != null)
            {
                m_blobSumRowsOfH.Dispose();
                m_blobSumRowsOfH = null;
            }
        }
 
        protected override void setup_internal_blobs(BlobCollection<T> col)
        {
            if (col.Count > 0)
                return;
 
            col.Add(m_blobInfoGain);
            col.Add(m_blobProb);
            col.Add(m_blobSumRowsOfH);
        }
 
        public override int ExactNumBottomBlobs
        {
            get { return -1; }
        }
 
        public override int MinBottomBlobs
        {
            get { return 2; }
        }
 
        public override int MaxBottomBlobs
        {
            get { return 3; }
        }
 
        public override int ExactNumTopBlobs
        {
            get { return -1; }
        }
 
        public override int MinTopBlobs
        {
            get { return 1; }
        }
 
        public override int MaxTopBlobs
        {
            get { return 2; }
        }
 
        protected virtual void sum_rows_of_H(Blob<T> blobH)
        {
            m_log.CHECK_EQ(blobH.count(), m_nNumLabels * m_nNumLabels, "H must be " + m_nNumLabels.ToString() + " x " + m_nNumLabels.ToString());
            float[] rgInfogainMat = convertF(blobH.update_cpu_data());
            float[] rgSum = convertF(m_blobSumRowsOfH.mutable_cpu_data);
 
            for (int row = 0; row < m_nNumLabels; row++)
            {
                rgSum[row] = 0;
 
                for (int col = 0; col < m_nNumLabels; col++)
                {
                    rgSum[row] += rgInfogainMat[row * m_nNumLabels + col];
                }
            }
 
            m_blobSumRowsOfH.mutable_cpu_data = convert(rgSum);
        }
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            base.LayerSetUp(colBottom, colTop);
 
            // Internal softmax layer.
            LayerParameter softmax_param = new LayerParameter(LayerParameter.LayerType.SOFTMAX);
            softmax_param.softmax_param.axis = m_param.infogain_loss_param.axis;
            softmax_param.loss_weight.Clear();
            softmax_param.loss_weight.Add(1);
            m_softmaxLayer = new SoftmaxLayer<T>(m_cuda, m_log, softmax_param);
            m_colSoftmaxBottomVec.Clear();
            m_colSoftmaxBottomVec.Add(colBottom[0]);
            m_colSoftmaxTopVec.Clear();
            m_colSoftmaxTopVec.Add(m_blobProb);
            m_softmaxLayer.Setup(m_colSoftmaxBottomVec, m_colSoftmaxTopVec);
 
            // ignore label.
            m_nIgnoreLabel = m_param.loss_param.ignore_label;
 
            // matrix H
            if (colBottom.Count < 3)
            {
                m_log.CHECK(m_param.infogain_loss_param.source != null, "Infogain matrix source must be specified.");
                PersistCaffe<T> persist = new PersistCaffe<T>(m_log, true);              
                BlobProto blobProto = persist.LoadBlobProto(m_param.infogain_loss_param.source, 1);
                m_blobInfoGain.FromProto(blobProto);
            }
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            base.Reshape(colBottom, colTop);
 
            m_softmaxLayer.Reshape(m_colSoftmaxBottomVec, m_colSoftmaxTopVec);
            m_nInfogainAxis = colBottom[0].CanonicalAxisIndex(m_param.infogain_loss_param.axis);
            m_nOuterNum = colBottom[0].count(0, m_nInfogainAxis);
            m_nInnerNum = colBottom[0].count(m_nInfogainAxis + 1);
 
            m_log.CHECK_EQ(m_nOuterNum * m_nInnerNum, colBottom[1].count(), "Number of labels must match the number of predictions; e.g., if infogain_axis == 1 and predictions shape is (N, C, H, W), label count (number of labels) must be N*H*W, with integer values in {0, 1, ..., C-1}.");
            m_nNumLabels = colBottom[0].shape(m_nInfogainAxis);
 
            Blob<T> blobInfoGain = null;
 
            if (colBottom.Count < 3)
                blobInfoGain = m_blobInfoGain;
            else
                blobInfoGain = colBottom[2];
 
            m_log.CHECK_EQ(blobInfoGain.count(), m_nNumLabels * m_nNumLabels, "The infogain count must equal 'num_labels' * 'num_labels'.");
            m_blobSumRowsOfH.Reshape(new List<int>() { m_nNumLabels });
            if (colBottom.Count == 2)
            {
                // H is provided as a parameter and will not change.  Sum rows once.
                sum_rows_of_H(blobInfoGain);
            }
            if (colTop.Count >= 2)
            {
                // softmax output.
                colTop[1].ReshapeLike(colBottom[0]);
            }
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            // The forward pass computes the softmax prob values.
            m_softmaxLayer.Forward(m_colSoftmaxBottomVec, m_colSoftmaxTopVec);
 
            Blob<T> blobInfoGain = (colBottom.Count < 3) ? m_blobInfoGain : colBottom[2];
            int nCount = 0;
            int nLabel = -1;
            double dfLoss = 0;
            double dfProb = 0;
            double dfProbLog = 0;
            double dfVal;
 
            if (typeof(T) == typeof(double))
            {
                double[] rgProbData = (double[])Convert.ChangeType(m_blobProb.update_cpu_data(), typeof(double[]));
                double[] rgBottomLabel = (double[])Convert.ChangeType(colBottom[1].update_cpu_data(), typeof(double[]));
                double[] rgInfoGainMat = (double[])Convert.ChangeType(blobInfoGain.update_cpu_data(), typeof(double[]));
 
                for (int i = 0; i < m_nOuterNum; i++)
                {
                    for (int j = 0; j < m_nInnerNum; j++)
                    {
                        nLabel = (int)rgBottomLabel[i * m_nInnerNum + j];
                        if (m_nIgnoreLabel.HasValue && m_nIgnoreLabel.Value == nLabel)
                            continue;
 
                        m_log.CHECK_GE(nLabel, 0, "The label should be greater than or equal to 0.");
                        m_log.CHECK_LT(nLabel, m_nNumLabels, "The label should be less than the number of labels '" + m_nNumLabels.ToString() + "'");
 
                        for (int l = 0; l < m_nNumLabels; l++)
                        {
                            dfProb = Math.Max(rgProbData[i * m_nInnerNum * m_nNumLabels + l * m_nInnerNum + j], kLOG_THRESHOLD);
                            dfProbLog = Math.Log(dfProb);
                            dfVal = rgInfoGainMat[nLabel * m_nNumLabels + l] * dfProbLog;
                            dfLoss -= dfVal;
                        }
 
                        nCount++;
                    }
                }
            }
            else
            {
                float[] rgProbData = (float[])Convert.ChangeType(m_blobProb.update_cpu_data(), typeof(float[]));
                float[] rgBottomLabel = (float[])Convert.ChangeType(colBottom[1].update_cpu_data(), typeof(float[]));
                float[] rgInfoGainMat = (float[])Convert.ChangeType(blobInfoGain.update_cpu_data(), typeof(float[]));
 
                for (int i = 0; i < m_nOuterNum; i++)
                {
                    for (int j = 0; j < m_nInnerNum; j++)
                    {
                        nLabel = (int)rgBottomLabel[i * m_nInnerNum + j];
                        if (m_nIgnoreLabel.HasValue && m_nIgnoreLabel.Value == nLabel)
                            continue;
 
                        m_log.CHECK_GE(nLabel, 0, "The label should be greater than or equal to 0.");
                        m_log.CHECK_LT(nLabel, m_nNumLabels, "The label should be less than the number of labels '" + m_nNumLabels.ToString() + "'");
 
                        for (int l = 0; l < m_nNumLabels; l++)
                        {
                            dfProb = Math.Max(rgProbData[i * m_nInnerNum * m_nNumLabels + l * m_nInnerNum + j], kLOG_THRESHOLD);
                            dfProbLog = Math.Log(dfProb);
                            dfVal = rgInfoGainMat[nLabel * m_nNumLabels + l] * dfProbLog;
                            dfLoss -= dfVal;
                        }
 
                        nCount++;
                    }
                }
            }
 
            double dfNormalizer = get_normalizer(m_normalization, nCount);
            double dfNormalizedLoss = dfLoss / dfNormalizer;
            colTop[0].SetData(dfNormalizedLoss, 0);
 
            if (colTop.Count == 2)
                colTop[1].ShareData(m_blobProb);
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            if (rgbPropagateDown[1])
                m_log.FAIL(type.ToString() + " Layer cannot backpropagate to label inputs.");
 
            if (rgbPropagateDown.Count > 2 && rgbPropagateDown[2])
                m_log.FAIL(type.ToString() + " Layer cannot backpropagate to infogain inputs.");
 
            if (rgbPropagateDown[0])
            {
                Blob<T> blobInfoGainMat = m_blobInfoGain;
                int nDim = colBottom[0].count() / m_nOuterNum;
                int nCount = 0;
                int nLabelValue = 0;
 
                if (colBottom.Count >= 3)
                {
                    blobInfoGainMat = colBottom[2];
                    // H is provided as a 'bottom' and might change so sum rows every time.
                    sum_rows_of_H(colBottom[2]);
                }
 
                if (typeof(T) == typeof(double))
                {
                    double[] rgProbData = (double[])Convert.ChangeType(m_blobProb.update_cpu_data(), typeof(double[]));
                    double[] rgBottomLabel = (double[])Convert.ChangeType(colBottom[1].update_cpu_data(), typeof(double[]));
                    double[] rgInfoGainMat = (double[])Convert.ChangeType(blobInfoGainMat.update_cpu_data(), typeof(double[]));
                    double[] rgSumsRowsH = (double[])Convert.ChangeType(m_blobSumRowsOfH.update_cpu_data(), typeof(double[]));
                    double[] rgBottomDiff = (double[])Convert.ChangeType(colBottom[0].mutable_cpu_diff, typeof(double[]));
 
                    for (int i=0; i<m_nOuterNum; i++)
                    {
                        for (int j=0; j<m_nInnerNum; j++)
                        {
                            nLabelValue = (int)rgBottomLabel[i * m_nInnerNum + j];
 
                            m_log.CHECK_GE(nLabelValue, 0, "The label should be greater than or equal to 0.");
                            m_log.CHECK_LT(nLabelValue, m_nNumLabels, "The label should be less than the number of labels '" + m_nNumLabels.ToString() + "'");
 
                            if (m_nIgnoreLabel.HasValue && m_nIgnoreLabel.Value == nLabelValue)
                            {
                                for (int l = 0; l < m_nNumLabels; l++)
                                {
                                    rgBottomDiff[i * nDim + l * m_nInnerNum + j] = 0;
                                }
                            }
                            else
                            {
                                for (int l = 0; l < m_nNumLabels; l++)
                                {
                                    int nIdx = i * nDim + l * m_nInnerNum + j;
                                    rgBottomDiff[nIdx] = rgProbData[nIdx] * rgSumsRowsH[nLabelValue] - rgInfoGainMat[nLabelValue * m_nNumLabels + l];
                                }
 
                                nCount++;
                            }
                        }
                    }
 
                    colBottom[0].mutable_cpu_diff = convert(rgBottomDiff);
                }
                else
                {
                    float[] rgProbData = (float[])Convert.ChangeType(m_blobProb.update_cpu_data(), typeof(float[]));
                    float[] rgBottomLabel = (float[])Convert.ChangeType(colBottom[1].update_cpu_data(), typeof(float[]));
                    float[] rgInfoGainMat = (float[])Convert.ChangeType(blobInfoGainMat.update_cpu_data(), typeof(float[]));
                    float[] rgSumsRowsH = (float[])Convert.ChangeType(m_blobSumRowsOfH.update_cpu_data(), typeof(float[]));
                    float[] rgBottomDiff = (float[])Convert.ChangeType(colBottom[0].mutable_cpu_diff, typeof(float[]));
 
                    for (int i = 0; i < m_nOuterNum; i++)
                    {
                        for (int j = 0; j < m_nInnerNum; j++)
                        {
                            nLabelValue = (int)rgBottomLabel[i * m_nInnerNum + j];
 
                            m_log.CHECK_GE(nLabelValue, 0, "The label should be greater than or equal to 0.");
                            m_log.CHECK_LT(nLabelValue, m_nNumLabels, "The label should be less than the number of labels '" + m_nNumLabels.ToString() + "'");
 
                            if (m_nIgnoreLabel.HasValue && m_nIgnoreLabel.Value == nLabelValue)
                            {
                                for (int l = 0; l < m_nNumLabels; l++)
                                {
                                    rgBottomDiff[i * nDim + l * m_nInnerNum + j] = 0;
                                }
                            }
                            else
                            {
                                for (int l = 0; l < m_nNumLabels; l++)
                                {
                                    int nIdx = i * nDim + l * m_nInnerNum + j;
                                    rgBottomDiff[nIdx] = rgProbData[nIdx] * rgSumsRowsH[nLabelValue] - rgInfoGainMat[nLabelValue * m_nNumLabels + l];
                                }
 
                                nCount++;
                            }
                        }
                    }
 
                    colBottom[0].mutable_cpu_diff = convert(rgBottomDiff);
                }
 
                // Scale the gradient
                double dfNormalizer = get_normalizer(m_normalization, nCount);
                double dfLossWeight = convertD(colTop[0].GetDiff(0)) / dfNormalizer;
                m_cuda.scal(colBottom[0].count(), dfLossWeight, colBottom[0].mutable_gpu_diff);
            }
        }
    }
}