OneHotLayer.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using MyCaffe.basecode;
using MyCaffe.common;
using MyCaffe.param;
using MyCaffe.fillers;
 
namespace MyCaffe.layers.nt
{
    public class OneHotLayer<T> : Layer<T>
    {
        int m_nAxis;
        BucketCollection m_colBuckets;
        float[] m_rgOneHotVector;
        float[] m_rgTop = null;
 
        public OneHotLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.ONEHOT;
        }
 
        protected override void dispose()
        {
            base.dispose();
        }
 
        public override int ExactNumBottomBlobs
        {
            get { return 1; }
        }
 
        public override int ExactNumTopBlobs
        {
            get { return 1; }
        }
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            m_rgOneHotVector = new float[m_param.onehot_param.num_output];
            m_colBuckets = new BucketCollection(m_param.onehot_param.min, m_param.onehot_param.max, (int)m_param.onehot_param.num_output);
            m_nAxis = colBottom[0].CanonicalAxisIndex(m_param.onehot_param.axis);
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            int nCount = colBottom[0].count(m_nAxis);
            m_log.CHECK_EQ(nCount, 1, "The bottom[0] count at axis " + m_nAxis.ToString() + " must equal 1");
 
            List<int> rgTopShape = Utility.Clone<int>(colBottom[0].shape());
            rgTopShape[m_nAxis] = m_colBuckets.Count;
 
            while (rgTopShape.Count < m_param.onehot_param.min_axes)
            {
                rgTopShape.Add(1);
            }
 
            colTop[0].Reshape(rgTopShape);
 
            int nTopCount = colTop[0].count();
            if (m_rgTop == null || m_rgTop.Length < nTopCount)
                m_rgTop = new float[nTopCount];
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            float[] rgBottom = convertF(colBottom[0].mutable_cpu_data);
            int nCount = colBottom[0].count(0, m_nAxis);
 
            for (int i = 0; i < nCount; i++)
            {
                int nIdx = m_colBuckets.Add(rgBottom[i]);
 
                for (int j = 0; j < m_rgOneHotVector.Length; j++)
                {
                    if (j == nIdx)
                        m_rgOneHotVector[j] = 1.0f;
                    else
                        m_rgOneHotVector[j] = 0;
                }
 
                Array.Copy(m_rgOneHotVector, 0, m_rgTop, i * m_rgOneHotVector.Length, m_rgOneHotVector.Length);
            }
 
            colTop[0].mutable_cpu_data = convert(m_rgTop);
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            int nItemCount = colTop[0].count(m_nAxis);
            m_log.CHECK_EQ(nItemCount, m_colBuckets.Count, "The count at the top[axis] is incorrect!");
 
            int nCount1 = colTop[0].count(0, m_nAxis);
            int nCount2 = colBottom[0].count(0, m_nAxis);
            m_log.CHECK_EQ(nCount1, nCount2, "The top and bottom have incompatible sizes.");
 
            // Convert top one-hot vectors to softmax indexes.
            float[] rgBottomDiff = convertF(colBottom[0].mutable_cpu_diff);
            float[] rgTopData = convertF(colTop[0].mutable_cpu_data);
            float[] rgTopDiff = convertF(colTop[0].mutable_cpu_diff);
 
            for (int i = 0; i < nCount1; i++)
            {
                int nItemIdx = i * nItemCount;
                float fDiff = 0;
                float fDiffSum = 0;
 
                for (int j = 0; j < nItemCount; j++)
                {
                    fDiff = rgTopDiff[nItemIdx + j];
 
                    if (rgTopData[nItemIdx + j] == 0)
                        fDiff *= -1;
 
                    fDiffSum += fDiff;
                }
 
                rgBottomDiff[i] = fDiffSum / nItemCount;
            }
 
            colBottom[0].mutable_cpu_diff = convert(rgBottomDiff);
        }
    }
}