ArgMaxLayer.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 ArgMaxLayer<T> : Layer<T>
    {
        bool m_bOutMaxVal;
        int m_nTopK;
        int? m_nAxis = null;
 
        public ArgMaxLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.ARGMAX;
        }
 
        public override int ExactNumBottomBlobs
        {
            get { return 1; }
        }
 
        public override int ExactNumTopBlobs
        {
            get { return 1; }
        }
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            ArgMaxParameter p = m_param.argmax_param;
            m_bOutMaxVal = p.out_max_val;
            m_nTopK = (int)p.top_k;
            m_nAxis = null;
 
            m_log.CHECK_GE(m_nTopK, 1, "Top k must not be less than 1.");
 
            if (p.axis.HasValue)
            {
                m_nAxis = colBottom[0].CanonicalAxisIndex(p.axis.Value);
                m_log.CHECK_GE(m_nAxis.Value, 0, "axis must not be less than zero.");
                m_log.CHECK_LE(m_nAxis.Value, colBottom[0].num_axes, "axis must be less tahn or equal to the dimension of the axis.");
                m_log.CHECK_LE(m_nTopK, colBottom[0].shape(m_nAxis.Value), "top_k must be less than or equal to the dimension of the axis.");
            }
            else
            {
                m_log.CHECK_LE(m_nTopK, colBottom[0].count(1), "top_k must be less than or equal to the dimension of the flattened bottom blob per instance.");
            }
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            int nNumTopAxes = colBottom[0].num_axes;
 
            if (nNumTopAxes < 3)
                nNumTopAxes = 3;
 
            List<int> rgShape = Utility.Create<int>(nNumTopAxes, 1);
 
            if (m_nAxis.HasValue)
            {
                // Produces max_ind or max_val per axis
                rgShape = Utility.Clone<int>(colBottom[0].shape());
                rgShape[m_nAxis.Value] = m_nTopK;
            }
            else
            {
                rgShape[0] = colBottom[0].shape(0);
                // Produces max_ind
                rgShape[2] = m_nTopK;
 
                if (m_bOutMaxVal)
                {
                    // Produces max_ind and max_val
                    rgShape[1] = 2;
                }
            }
 
            colTop[0].Reshape(rgShape);
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            if (m_param.argmax_param.enable_cuda_impl)
                forward_gpu(colBottom, colTop);
            else
                forward_cpu(colBottom, colTop);
        }
 
        private void forward_gpu(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            int nAxis = m_nAxis.GetValueOrDefault(1);
            int nOuterNum = colBottom[0].count(0, nAxis);
            int nChannels = colBottom[0].count(nAxis);
            int nInnerNum = 1;
 
            if (m_nTopK > 1)
                m_log.WriteLine("WARNING: The gpu implementation of argmax only supports TopK = 1.");
 
            if (m_bOutMaxVal)
                m_log.WriteLine("WARNING: Currently the gpu implementation of argmax does now support output of the max values.");
 
            if (m_param.argmax_param.operation == ArgMaxParameter.COMPARE_OPERATOR.MAX)
                m_cuda.channel_max(colBottom[0].count(), nOuterNum, nChannels, nInnerNum, colBottom[0].gpu_data, colTop[0].mutable_gpu_data, true);
            else
                m_cuda.channel_min(colBottom[0].count(), nOuterNum, nChannels, nInnerNum, colBottom[0].gpu_data, colTop[0].mutable_gpu_data, true);
        }
 
        private void forward_cpu(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            double[] rgBottomData = convertD(colBottom[0].update_cpu_data());
            double[] rgTopData = convertD(colTop[0].mutable_cpu_data);
            int nDim;
            int nAxisDist;
 
            if (m_nAxis.HasValue)
            {
                nDim = colBottom[0].shape(m_nAxis.Value);
                // Distance between values of axis in blob.
                nAxisDist = colBottom[0].count(m_nAxis.Value) / nDim;
            }
            else
            {
                nDim = colBottom[0].count(1);
                nAxisDist = 1;
            }
 
            int nNum = colBottom[0].count() / nDim;
 
            for (int i = 0; i < nNum; i++)
            {
                List<KeyValuePair<double, int>> rgBottomDataPair = new List<KeyValuePair<double, int>>();
 
                for (int j = 0; j < nDim; j++)
                {
                    int nIdx = (i / nAxisDist * nDim + j) * nAxisDist + i % nAxisDist;
                    rgBottomDataPair.Add(new KeyValuePair<double, int>(rgBottomData[nIdx], j));
                }
 
                if (m_param.argmax_param.operation == ArgMaxParameter.COMPARE_OPERATOR.MIN)
                    rgBottomDataPair.Sort(new Comparison<KeyValuePair<double, int>>(sortDataItemsDescending));
                else
                    rgBottomDataPair.Sort(new Comparison<KeyValuePair<double, int>>(sortDataItemsAscending));
 
                for (int j = 0; j < m_nTopK; j++)
                {
                    if (m_bOutMaxVal)
                    {
                        if (m_nAxis.HasValue)
                        {
                            // Produces max_val per axis
                            int nIdx = (i / nAxisDist * m_nTopK + j) * nAxisDist + i % nAxisDist;
                            rgTopData[nAxisDist] = rgBottomDataPair[j].Key;
                        }
                        else
                        {
                            // Produces max_ind and max_val
                            int nIdx1 = 2 * i * m_nTopK + j;
                            rgTopData[nIdx1] = rgBottomDataPair[j].Value;
                            int nIdx2 = 2 * i * m_nTopK + m_nTopK + j;
                            rgTopData[nIdx2] = rgBottomDataPair[j].Key;
                        }
                    }
                    else
                    {
                        // Produces max_ind per axis.
                        int nIdx = (i / nAxisDist * m_nTopK + j) * nAxisDist + i % nAxisDist;
                        rgTopData[nIdx] = rgBottomDataPair[j].Value;
                    }
                }
            }
 
            colTop[0].mutable_cpu_data = convert(rgTopData);
        }
 
        private int sortDataItemsAscending(KeyValuePair<double, int> a, KeyValuePair<double, int> b)
        {
            if (a.Key < b.Key)
                return 1;
 
            if (a.Key > b.Key)
                return -1;
 
            return 0;
        }
 
        private int sortDataItemsDescending(KeyValuePair<double, int> a, KeyValuePair<double, int> b)
        {
            if (a.Key < b.Key)
                return -1;
 
            if (a.Key > b.Key)
                return 1;
 
            return 0;
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            new NotImplementedException();
        }
    }
}