MVNLayer.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 MVNLayer<T> : Layer<T>
    {
        Blob<T> m_blobMean;
        Blob<T> m_blobVariance;
        Blob<T> m_blobTemp;
        // Sum_multiplier is used to carry out sum using BLAS
        Blob<T> m_blobSumMultiplier;
        double m_dfEps;
 
        public MVNLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.MVN;
            m_blobMean = new common.Blob<T>(cuda, log);
            m_blobMean.Name = m_param.name + " mean";
            m_blobVariance = new common.Blob<T>(cuda, log);
            m_blobVariance.Name = m_param.name + " variance";
            m_blobTemp = new Blob<T>(cuda, log);
            m_blobTemp.Name = m_param.name + " temp";
            m_blobSumMultiplier = new Blob<T>(cuda, log);
            m_blobSumMultiplier.Name = m_param.name + " summult";
        }
 
        protected override void dispose()
        {
            if (m_blobMean != null)
            {
                m_blobMean.Dispose();
                m_blobMean = null;
            }
 
            if (m_blobVariance != null)
            {
                m_blobVariance.Dispose();
                m_blobVariance = null;
            }
 
            if (m_blobTemp != null)
            {
                m_blobTemp.Dispose();
                m_blobTemp = null;
            }
 
            if (m_blobSumMultiplier != null)
            {
                m_blobSumMultiplier.Dispose();
                m_blobSumMultiplier = null;
            }
 
            base.dispose();
        }
 
        protected override void setup_internal_blobs(BlobCollection<T> col)
        {
            if (col.Count > 0)
                return;
 
            col.Add(m_blobMean);
            col.Add(m_blobVariance);
            col.Add(m_blobTemp);
            col.Add(m_blobSumMultiplier);
        }
 
        public override int ExactNumBottomBlobs
        {
            get { return 1; }
        }
 
        public override int ExactNumTopBlobs
        {
            get { return 1; }
        }
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            colTop[0].ReshapeLike(colBottom[0]);
 
            m_blobMean.Reshape(colBottom[0].num, colBottom[0].channels, 1, 1);
            m_blobVariance.Reshape(colBottom[0].num, colBottom[0].channels, 1, 1);
            m_blobTemp.Reshape(colBottom[0].num, colBottom[0].channels, colBottom[0].height, colBottom[0].width);
 
            if (m_param.mvn_param.across_channels)
                m_blobSumMultiplier.Reshape(1, colBottom[0].channels, colBottom[0].height, colBottom[0].width);
            else
                m_blobSumMultiplier.Reshape(1, 1, colBottom[0].height, colBottom[0].width);
 
            m_blobSumMultiplier.SetData(1.0);
            m_dfEps = m_param.mvn_param.eps;
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            long hBottomData = colBottom[0].gpu_data;
            long hTopData = colTop[0].mutable_gpu_data;
            int nNum = colBottom[0].num;
 
            if (!m_param.mvn_param.across_channels)
                nNum *= colBottom[0].channels;
 
            int nDim = colBottom[0].count() / nNum;
 
            // subtract mean
            m_cuda.gemv(false, nNum, nDim, 1.0 / nDim, hBottomData, m_blobSumMultiplier.gpu_data, 0.0, m_blobMean.mutable_gpu_data); //EX
            m_cuda.gemm(false, false, nNum, nDim, 1, -1.0, m_blobMean.gpu_data, m_blobSumMultiplier.gpu_data, 0.0, m_blobTemp.mutable_gpu_data);
            m_cuda.add(m_blobTemp.count(), hBottomData, m_blobTemp.gpu_data, hTopData); // X-EX
 
            if (m_param.mvn_param.normalize_variance)
            {
                // compute variance using var(X) = E((X-EX)^2)
                m_cuda.powx(colBottom[0].count(), hTopData, 2.0, m_blobTemp.mutable_gpu_data); // (X-EX)^2
                m_cuda.gemv(false, nNum, nDim, 1.0 / nDim, m_blobTemp.gpu_data, m_blobSumMultiplier.gpu_data, 0.0, m_blobVariance.mutable_gpu_data); // E((X-EX)^2)
 
                // normalize variance
                m_cuda.powx(m_blobVariance.count(), m_blobVariance.gpu_data, 0.5, m_blobVariance.mutable_gpu_data);
                m_cuda.add_scalar(m_blobVariance.count(), m_dfEps, m_blobVariance.mutable_gpu_data);
                m_cuda.gemm(false, false, nNum, nDim, 1, 1.0, m_blobVariance.gpu_data, m_blobSumMultiplier.gpu_data, 0.0, m_blobTemp.mutable_gpu_data);
                m_cuda.div(m_blobTemp.count(), hTopData, m_blobTemp.gpu_data, hTopData);
            }
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            if (!rgbPropagateDown[0])
                return;
 
            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;
            int nNum = colBottom[0].num;
 
            if (!m_param.mvn_param.across_channels)
                nNum *= colBottom[0].channels;
 
            int nDim = colBottom[0].count() / nNum;
 
            if (m_param.mvn_param.normalize_variance)
            {
                m_cuda.mul(m_blobTemp.count(), hTopData, hTopDiff, hBottomDiff);
                m_cuda.gemv(false, nNum, nDim, 1.0, hBottomDiff, m_blobSumMultiplier.gpu_data, 0.0, m_blobMean.mutable_gpu_data);
                m_cuda.gemm(false, false, nNum, nDim, 1, 1.0, m_blobMean.gpu_data, m_blobSumMultiplier.gpu_data, 0.0, hBottomDiff);
                m_cuda.mul(m_blobTemp.count(), hTopData, hBottomDiff, hBottomDiff);
 
                m_cuda.gemv(false, nNum, nDim, 1.0, hTopDiff, m_blobSumMultiplier.gpu_data, 0.0, m_blobMean.mutable_gpu_data);
                m_cuda.gemm(false, false, nNum, nDim, 1, 1.0, m_blobMean.gpu_data, m_blobSumMultiplier.gpu_data, 1.0, hBottomDiff);
 
                m_cuda.axpby(m_blobTemp.count(), 1.0, hTopDiff, -1.0 / nDim, hBottomDiff);
 
                // put the squares of bottom into temp_
                m_cuda.powx(m_blobTemp.count(), hBottomData, 2.0, m_blobTemp.mutable_gpu_data);
                m_cuda.gemm(false, false, nNum, nDim, 1, 1.0, m_blobVariance.gpu_data, m_blobSumMultiplier.gpu_data, 0.0, m_blobTemp.mutable_gpu_data);
                m_cuda.div(m_blobTemp.count(), hBottomDiff, m_blobTemp.gpu_data, hBottomDiff);
            }
            else
            {
                m_cuda.gemv(false, nNum, nDim, 1.0 / nDim, hTopDiff, m_blobSumMultiplier.gpu_data, 0.0, m_blobMean.mutable_gpu_data);
                m_cuda.gemm(false, false, nNum, nDim, 1, -1.0, m_blobMean.gpu_data, m_blobSumMultiplier.gpu_data, 0.0, m_blobTemp.mutable_gpu_data);
                m_cuda.add(m_blobTemp.count(), hTopDiff, m_blobTemp.gpu_data, hBottomDiff);
            }
        }
    }
}