MemoryLossLayer.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 MemoryLossLayer<T> : LossLayer<T>
    {
        object m_userState = null;
        bool m_bWarningMade = false;
        bool m_bEnableLoss = true;
 
        public event EventHandler<MemoryLossLayerGetLossArgs<T>> OnGetLoss;
 
        public MemoryLossLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.MEMORY_LOSS;
        }
 
        protected override void dispose()
        {
            base.dispose();
        }
 
        public object user_state
        {
            get { return m_userState; }
            set { m_userState = value; }
        }
 
        public override int ExactNumBottomBlobs
        {
            get { return -1; }
        }
 
        public override int MinBottomBlobs
        {
            get { return 1; }
        }
 
        public override int MaxBottomBlobs
        {
            get { return 2; }
        }
 
        public override int ExactNumTopBlobs
        {
            get { return 1; }
        }
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            base.LayerSetUp(colBottom, colTop);
 
            if (!m_param.loss_param.normalization.HasValue)
                m_normalization = (m_param.loss_param.normalize) ? LossParameter.NormalizationMode.VALID : LossParameter.NormalizationMode.BATCH_SIZE;
            else
                m_normalization = m_param.loss_param.normalization.Value;
 
            m_bWarningMade = false;
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            bool bUniformSize = true;
            int nAxis = colBottom[0].CanonicalAxisIndex(1);
            int nCount = colBottom[0].count(nAxis);
 
            for (int i = 1; i < colBottom.Count; i++)
            {
                int nCount1 = colBottom[i].count(nAxis);
                if (nCount1 != nCount)
                {
                    bUniformSize = false;
                    break;
                }
            }
 
            if (!bUniformSize)
            {
                if (!m_bWarningMade)
                {
                    m_log.WriteLine("WARNING: The MemoryDataLayer bottoms are not of uniform size, so the normalization will be set to NONE.");
                    m_bWarningMade = true;
                }
 
                m_normalization = LossParameter.NormalizationMode.NONE;
                m_nOuterNum = 0;
                m_nInnerNum = 0;
            }
            else
            {
                m_nOuterNum = colBottom[0].count(0, nAxis);
                m_nInnerNum = colBottom[0].count(nAxis + 1);
            }
 
            List<int> rgLossShape = new List<int>(); // Loss layers output a scalar, 0 axes.
            colTop[0].Reshape(rgLossShape);
            colTop[0].type = BLOB_TYPE.LOSS;
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            if (OnGetLoss == null)
                m_log.FAIL("The OnGetLoss event must be implemented.  Make sure the SolverParameter 'custom_trainer' points to a trainer that connects the OnGetLoss event.");
 
            double dfNormalizer = get_normalizer(m_normalization, -1);
            MemoryLossLayerGetLossArgs<T> e = new MemoryLossLayerGetLossArgs<T>(colBottom, m_userState, dfNormalizer);
            OnGetLoss(this, e);
 
            m_bEnableLoss = e.EnableLossUpdate;
            colTop[0].SetData(e.Loss / dfNormalizer, 0);
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            if (!rgbPropagateDown[0])
                return;
 
            double dfNormalizer = get_normalizer(m_normalization, -1);
 
            // mutliply the loss by the loss weight (in top[0].diff)
            if (m_bEnableLoss)
            {
                double dfLoss = convertD(colTop[0].GetData(0));
 
                for (int i = 0; i < colBottom.Count; i++)
                {
                    m_cuda.copy(colBottom[i].count(), colBottom[i].gpu_data, colBottom[i].mutable_gpu_diff);
                    m_cuda.mul_scalar(colBottom[i].count(), dfLoss, colBottom[i].mutable_gpu_diff);
                }
 
                dfNormalizer = 1.0;
            }
 
            double dfTopDiff = convertD(colTop[0].GetDiff(0)); // loss weight
            double dfLossWeight = dfTopDiff / dfNormalizer;
 
            // Apply the loss weight to the bottom diffs.
            if (dfLossWeight != 1.0)
            {
                for (int i = 0; i < colBottom.Count; i++)
                {
                    m_cuda.scal(colBottom[i].count(), convert(dfLossWeight), colBottom[i].mutable_gpu_diff);
                }
            }
        }
    }
 
    public class MemoryLossLayerGetLossArgs<T> : EventArgs
    {
        object m_userState = null;
        double m_dfLoss = 0;
        double m_dfNormalizer = 1;
        BlobCollection<T> m_colBottom;
        bool m_bEnableLossUpdate = true;
        object m_tag = null;
 
        public MemoryLossLayerGetLossArgs(BlobCollection<T> colBottom, object userState, double dfNormalizer)
        {
            m_userState = userState;
            m_colBottom = colBottom;
        }
 
        public object user_state
        {
            get { return m_userState; }
        }
 
        public BlobCollection<T> Bottom
        {
            get { return m_colBottom; }
        }
 
        public double Normalizer
        {
            get { return m_dfNormalizer; }
        }
 
        public double Loss
        {
            get { return m_dfLoss; }
            set { m_dfLoss = value; }
        }
 
        public bool EnableLossUpdate
        {
            get { return m_bEnableLossUpdate; }
            set { m_bEnableLossUpdate = value; }
        }
 
        public object Tag
        {
            get { return m_tag; }
            set { m_tag = value; }
        }
    }
}