Normalization2Layer.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using MyCaffe.basecode;
using MyCaffe.common;
using MyCaffe.fillers;
using MyCaffe.param;
using MyCaffe.param.ssd;
 
namespace MyCaffe.layers.ssd
{
    public class Normalization2Layer<T> : Layer<T>
    {
        Blob<T> m_blobNorm;
        Blob<T> m_blobSumChannelMultiplier;
        Blob<T> m_blobSumSpatialMultiplier;
        Blob<T> m_blobBuffer;
        Blob<T> m_blobBufferChannel;
        Blob<T> m_blobBufferSpatial;
        bool m_bAcrossSpatial;
        bool m_bChannelShared;
        float m_fEps;
 
        public Normalization2Layer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.NORMALIZATION2;
            m_blobNorm = new Blob<T>(cuda, log, false);
            m_blobNorm.Name = m_param.name + " norm";
            m_blobSumChannelMultiplier = new Blob<T>(cuda, log, false);
            m_blobSumChannelMultiplier.Name = m_param.name + " sum_chan_mult";
            m_blobSumSpatialMultiplier = new Blob<T>(cuda, log, false);
            m_blobSumSpatialMultiplier.Name = m_param.name + " sum_spat_mult";
            m_blobBuffer = new Blob<T>(cuda, log, false);
            m_blobBuffer.Name = m_param.name + " buffer";
            m_blobBufferChannel = new Blob<T>(cuda, log, false);
            m_blobBufferChannel.Name = m_param.name + " buffer_chan";
            m_blobBufferSpatial = new Blob<T>(cuda, log, false);
            m_blobBufferSpatial.Name = m_param.name + " buffer_spat";
        }
 
        protected override void dispose()
        {
            m_blobNorm.Dispose();
            m_blobSumChannelMultiplier.Dispose();
            m_blobSumSpatialMultiplier.Dispose();
            m_blobBuffer.Dispose();
            m_blobBufferChannel.Dispose();
            m_blobBufferSpatial.Dispose();
            base.dispose();
        }
 
        protected override void setup_internal_blobs(BlobCollection<T> col)
        {
            if (col.Count > 0)
                return;
 
            col.Add(m_blobNorm);
            col.Add(m_blobSumChannelMultiplier);
            col.Add(m_blobSumSpatialMultiplier);
            col.Add(m_blobBuffer);
            col.Add(m_blobBufferChannel);
            col.Add(m_blobBufferSpatial);
        }
 
        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_log.CHECK_GE(colBottom[0].num_axes, 2, "The bottom(0) must have >= 2 axes.");
            m_blobBuffer.Reshape(1, colBottom[0].channels, colBottom[0].height, colBottom[0].width);
            m_blobBufferChannel.Reshape(1, colBottom[0].channels, 1, 1);
            m_blobBufferSpatial.Reshape(1, 1, colBottom[0].height, colBottom[0].width);
 
            Normalization2Parameter norm_param = layer_param.normalization2_param;
            m_bAcrossSpatial = norm_param.across_spatial;
 
            if (m_bAcrossSpatial)
                m_blobNorm.Reshape(colBottom[0].num, 1, 1, 1);
            else
                m_blobNorm.Reshape(colBottom[0].num, 1, colBottom[0].height, colBottom[0].width);
 
            m_fEps = norm_param.eps;
 
            int nChannels = colBottom[0].channels;
            int nSpatialDim = colBottom[0].width * colBottom[0].height;
 
            m_blobSumChannelMultiplier.Reshape(1, nChannels, 1, 1);
            m_blobSumChannelMultiplier.SetData(1);
 
            m_blobSumSpatialMultiplier.Reshape(1, 1, colBottom[0].height, colBottom[0].width);
            m_blobSumSpatialMultiplier.SetData(1);
 
            m_bChannelShared = norm_param.channel_shared;
 
            if (m_colBlobs.Count > 0)
            {
                m_log.WriteLine("Skipping parameter initialization.");
            }
            else
            {
                Blob<T> blobScale = new Blob<T>(m_cuda, m_log);
                blobScale.Name = m_param.name + " scale";
                List<int> rgShape = new List<int>();
 
                if (!m_bChannelShared)
                    rgShape.Add(nChannels);
 
                blobScale.Reshape(rgShape);
                Filler<T> filler = Filler<T>.Create(m_cuda, m_log, norm_param.scale_filler);
                filler.Fill(blobScale);
 
                m_colBlobs.Add(blobScale);
            }
 
            if (m_bChannelShared)
                m_log.CHECK_EQ(m_colBlobs[0].count(), 1, "Scale size is inconsistent with prototxt config.");
            else
                m_log.CHECK_EQ(m_colBlobs[0].count(), nChannels, "Scale size is inconsistent with prototxt config.");
 
            m_rgbParamPropagateDown[0] = true;
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            m_log.CHECK_GE(colBottom[0].num_axes, 2, "Number of axes of bottom blob must be >= 2.");
            colTop[0].ReshapeLike(colBottom[0]);
            m_blobBuffer.Reshape(1, colBottom[0].channels, colBottom[0].height, colBottom[0].width);
 
            if (!m_bAcrossSpatial)
                m_blobNorm.Reshape(colBottom[0].num, 1, colBottom[0].height, colBottom[0].width);
 
            int nSpatialDim = colBottom[0].height * colBottom[0].width;
 
            if (nSpatialDim != m_blobSumSpatialMultiplier.count())
            {
                m_blobSumSpatialMultiplier.Reshape(1, 1, colBottom[0].height, colBottom[0].width);
                m_blobSumSpatialMultiplier.SetData(1);
                m_blobBufferSpatial.ReshapeLike(m_blobSumSpatialMultiplier);
            }
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            long hBottomData = colBottom[0].gpu_data;
            long hTopData = colTop[0].mutable_gpu_data;
            long hBufferData = m_blobBuffer.mutable_gpu_data;
            long hScale = 0;
            long hSumChannelMultiplier = m_blobSumChannelMultiplier.gpu_data;
            long hNormData = 0;
            T[] rgNormData = null;
            T[] rgScale = null;
 
            if (m_bAcrossSpatial)
            {
                // Need to index it.
                rgNormData = m_blobNorm.mutable_cpu_data;
            }
            else
            {
                hNormData = m_blobNorm.mutable_gpu_data;
                // Add eps to avoid overflow.
                m_blobNorm.SetData(m_fEps);
            }
 
            if (m_bChannelShared)
                rgScale = m_colBlobs[0].mutable_cpu_data;
            else
                hScale = m_colBlobs[0].gpu_data;
 
            int nNum = colBottom[0].num;
            int nDim = colBottom[0].count() / nNum;
            int nSpatialDim = colBottom[0].height * colBottom[0].width;
            int nChannels = colBottom[0].channels;
            int nNormDataOffset = 0;
            int nBottomOffset = 0;
            int nTopOffset = 0;
 
            for (int n = 0; n < nNum; n++)
            {
                m_cuda.powx(nDim, hBottomData, 2.0, hBufferData, nBottomOffset, 0);
 
                if (m_bAcrossSpatial)
                {
                    double dfNormSqr = Utility.ConvertVal<T>(m_blobBuffer.asum_data());
                    dfNormSqr += m_fEps; // Add eps to avoid overflow.
                    double dfNorm = Math.Pow(dfNormSqr, 0.5);
                    m_cuda.scale(nDim, Utility.ConvertVal<T>(1.0 / dfNorm), hBottomData, hTopData, nBottomOffset, nTopOffset);
                    rgNormData[n] = Utility.ConvertVal<T>(dfNorm);
                }
                else
                {
                    // compute norm
                    m_cuda.gemv(true, nChannels, nSpatialDim, m_tOne, hBufferData, hSumChannelMultiplier, m_tOne, hNormData, 0, 0, nNormDataOffset);
                    m_cuda.powx(nSpatialDim, hNormData, 0.5, hNormData, nNormDataOffset, nNormDataOffset);
 
                    // Scale the layer.
                    m_cuda.divbsx(nDim, hBottomData, nBottomOffset, hNormData, nNormDataOffset, nChannels, nSpatialDim, false, hTopData, nTopOffset);
                    nNormDataOffset += nSpatialDim;
                }
 
                // Scale the output
                if (m_bChannelShared)
                {
                    m_cuda.scal(nDim, rgScale[0], hTopData, nTopOffset);
                }
                else
                {
                    m_cuda.mulbsx(nDim, hTopData, nTopOffset, hScale, 0, nChannels, nSpatialDim, true, hTopData, nTopOffset);
                }
 
                nBottomOffset += nDim;
                nTopOffset += nDim;
            }
 
            if (rgNormData != null)
                m_blobNorm.mutable_cpu_data = rgNormData;
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            long hTopDiff = colTop[0].gpu_diff;
            long hTopData = colTop[0].gpu_data;
            long hBottomData = colBottom[0].gpu_data;
            long hBottomDiff = colBottom[0].mutable_gpu_diff;
            long hBufferData = m_blobBuffer.mutable_gpu_data;
            long hScale = 0;
            long hBufferChannel = m_blobBufferChannel.mutable_gpu_data;
            long hBufferSpatial = m_blobBufferSpatial.mutable_gpu_data;
            long hSumChannelMultiplier = m_blobSumChannelMultiplier.gpu_data;
            long hSumSpatialMultiplier = m_blobSumSpatialMultiplier.gpu_data;
            long hNormData = 0;
            double dfScale = 0;
            double[] rgNormData = null;
 
            if (m_bAcrossSpatial)
            {
                // Need to index it.
                rgNormData = Utility.ConvertVec<T>(m_blobNorm.mutable_cpu_data);
            }
            else
            {
                hNormData = m_blobNorm.mutable_gpu_data;
            }
 
            if (m_bChannelShared)
                dfScale = Utility.ConvertVal<T>(m_colBlobs[0].GetData(0));
            else
                hScale = m_colBlobs[0].mutable_gpu_data;
 
            int nCount = colTop[0].count();
            int nNum = colTop[0].num;
            int nDim = nCount / nNum;
            int nSpatialDim = colTop[0].height * colTop[0].width;
            int nChannels = colTop[0].channels;
 
            // Propagate to param
            if (m_rgbParamPropagateDown[0])
            {
                if (m_bChannelShared)
                {
                    double dfScaleDiff = Utility.ConvertVal<T>(m_colBlobs[0].GetDiff(0));
                    double dfA = Utility.ConvertVal<T>(m_cuda.dot(nCount, hTopData, hTopDiff));
 
                    dfScaleDiff += (dfA / dfScale);
                    m_colBlobs[0].SetDiff(dfScaleDiff, 0);
                }
                else
                {
                    long hScaleDiff = m_colBlobs[0].mutable_gpu_diff;
 
                    for (int n = 0; n < nNum; n++)
                    {
                        // Compute A
                        m_cuda.mul(nDim, hTopData, hTopDiff, hBufferData, n * nDim, n * nDim, 0);
                        m_cuda.gemv(false, nChannels, nSpatialDim, m_tOne, hBufferData, hSumSpatialMultiplier, m_tZero, hBufferChannel);
 
                        // Store A / scale[i] in bufferdata temporarily.
                        m_cuda.div(nChannels, hBufferChannel, hScale, hBufferChannel);
                        m_cuda.add(nChannels, hBufferChannel, hScaleDiff, hScaleDiff);
                    }
                }
            }
 
            // Propagate to bottom
            if (rgbPropagateDown[0])
            {
                int nBottomDataOffset = 0;
                int nBottomDiffOffset = 0;
                int nTopDiffOffset = 0;
                int nNormDataOffset = 0;
 
                for (int n = 0; n < nNum; n++)
                {
                    if (m_bAcrossSpatial)
                    {
                        double dfA = Utility.ConvertVal<T>(m_cuda.dot(nDim, hBottomData, hTopDiff, nBottomDataOffset, nTopDiffOffset));
                        double dfScale1 = dfA / rgNormData[n] / rgNormData[n];
                        m_cuda.scale(nDim, Utility.ConvertVal<T>(dfScale1), hBottomData, hBottomDiff, nBottomDataOffset, nBottomDiffOffset);
                        m_cuda.sub(nDim, hTopDiff, hBottomDiff, hBottomDiff, nTopDiffOffset, nBottomDiffOffset, nBottomDiffOffset);
 
                        dfScale1 = 1.0 / rgNormData[n];
                        m_cuda.scale(nDim, Utility.ConvertVal<T>(dfScale1), hBottomDiff, hBottomDiff, nBottomDiffOffset, nBottomDiffOffset);
                    }
                    else
                    {
                        // dot product between bottom_data and top_diff
                        m_cuda.mul(nDim, hBottomData, hTopDiff, hBufferData, nBottomDataOffset, nTopDiffOffset, 0);
                        m_cuda.gemv(true, nChannels, nSpatialDim, m_tOne, hBufferData, hSumChannelMultiplier, m_tZero, hBufferSpatial);
 
                        // scale bottom diff
                        m_cuda.mulbsx(nDim, hBottomData, nBottomDataOffset, hBufferSpatial, 0, nChannels, nSpatialDim, false, hBottomDiff, nBottomDiffOffset);
 
                        // divide by square of norm
                        m_cuda.powx(nSpatialDim, hNormData, 2.0, hBufferSpatial);
                        m_cuda.divbsx(nDim, hBottomDiff, nBottomDiffOffset, hBufferSpatial, 0, nChannels, nSpatialDim, false, hBottomDiff, nBottomDiffOffset);
                        m_cuda.sub(nDim, hTopDiff, hBottomDiff, hBottomDiff, nTopDiffOffset, nBottomDiffOffset, nBottomDiffOffset);
 
                        // divide by norm
                        m_cuda.divbsx(nDim, hBottomDiff, nBottomDiffOffset, hNormData, nNormDataOffset, nChannels, nSpatialDim, false, hBottomDiff, nBottomDiffOffset);
                        nNormDataOffset += nSpatialDim;
                    }
 
                    // Scale the diff
                    if (m_bChannelShared)
                        m_cuda.scal(nDim, dfScale, hBottomDiff, nBottomDiffOffset);
                    else
                        m_cuda.mulbsx(nDim, hBottomDiff, nBottomDiffOffset, hScale, 0, nChannels, nSpatialDim, true, hBottomDiff, nBottomDiffOffset);
 
                    nBottomDataOffset += nDim;
                    nBottomDiffOffset += nDim;
                    nTopDiffOffset += nDim;
                }
            }
        }
    }
}