ScaleLayer.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
{
    public class ScaleLayer<T> : Layer<T>
    {
        BiasLayer<T> m_biasLayer = null;
        BlobCollection<T> m_colBiasBottomVec = new BlobCollection<T>();
        List<bool> m_rgbBiasPropagateDown = new List<bool>();
        int m_nBiasParamId;
        Blob<T> m_blobSumMultiplier;
        Blob<T> m_blobSumResult;
        Blob<T> m_blobTemp;
        int m_nAxis;
        int m_nOuterDim;
        int m_nScaleDim;
        int m_nInnerDim;
 
        public ScaleLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.SCALE;
            m_blobSumMultiplier = new Blob<T>(cuda, log);
            m_blobSumMultiplier.Name = m_param.name + " summult";
            m_blobSumResult = new Blob<T>(cuda, log);
            m_blobSumResult.Name = m_param.name + " sumres";
            m_blobTemp = new Blob<T>(cuda, log);
            m_blobTemp.Name = m_param.name + " temp";
        }
 
        protected override void dispose()
        {
            if (m_biasLayer != null)
            {
                m_biasLayer.Dispose();
                m_biasLayer = null;
            }
 
            if (m_blobSumMultiplier != null)
            {
                m_blobSumMultiplier.Dispose();
                m_blobSumMultiplier = null;
            }
 
            if (m_blobSumResult != null)
            {
                m_blobSumResult.Dispose();
                m_blobSumResult = null;
            }
 
            if (m_blobTemp != null)
            {
                m_blobTemp.Dispose();
                m_blobTemp = null;
            }
 
            base.dispose();
        }
 
        protected override void setup_internal_blobs(BlobCollection<T> col)
        {
            if (col.Count > 0)
                return;
 
            col.Add(m_blobSumMultiplier);
            col.Add(m_blobSumResult);
            col.Add(m_blobTemp);
        }
 
        public override int MinBottomBlobs
        {
            get { return 1; }
        }
 
        public override int MaxBottomBlobs
        {
            get { return 2; }
        }
 
        public override int ExactNumTopBlobs
        {
            get { return 1; }
        }
 
        public override bool ReInitializeParameters(WEIGHT_TARGET target)
        {
            base.ReInitializeParameters(target);
 
            if (target == WEIGHT_TARGET.BOTH || target == WEIGHT_TARGET.WEIGHTS)
            {
                FillerParameter fp = m_param.scale_param.filler;
                if (fp == null)
                    fp = new FillerParameter("constant", 1.0);
 
                Filler<T> filler = Filler<T>.Create(m_cuda, m_log, fp);
                filler.Fill(m_colBlobs[0]);
            }
 
            if (m_param.scale_param.bias_term)
                m_biasLayer.ReInitializeParameters(target);
 
            return true;
        }
 
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            ScaleParameter p = m_param.scale_param;
 
            if (colBottom.Count == 1 && blobs.Count > 0)
            {
                m_log.WriteLine("Skipping parameter initialization.");
            }
            else if (colBottom.Count == 1)
            {
                // scale is a learned parameter; initialize it.
                m_nAxis = colBottom[0].CanonicalAxisIndex(p.axis);
                int nNumAxes = p.num_axes;
                m_log.CHECK_GE(nNumAxes, -1, "num_axes must be non-negative, or -1 to extend to the end of bottom[0].");
 
                if (nNumAxes >= 0)
                    m_log.CHECK_GE(colBottom[0].num_axes, m_nAxis + nNumAxes, "scale blob's shape extends past bottom[0]'s shape when applied starting with bottom[0] axis = " + m_nAxis.ToString());
 
                m_colBlobs = new BlobCollection<T>();
 
                List<int> rgShape = new List<int>();
                int nStart = m_nAxis;
                int nEnd = (nNumAxes == -1) ? colBottom[0].shape().Count : nStart + nNumAxes;
 
                for (int i = nStart; i < nEnd; i++)
                {
                    rgShape.Add(colBottom[0].shape(i));
                }
 
                Blob<T> blobScale = new Blob<T>(m_cuda, m_log);
                blobScale.Name = m_param.name + " scale";
                blobScale.type = BLOB_TYPE.INTERNAL;
 
                if (!shareParameter(blobScale, rgShape))
                {
                    blobScale.Reshape(rgShape);
                    FillerParameter fp = p.filler;
 
                    // Default to unit (1) filler for identity operation.
                    if (fp == null)
                        fp = new FillerParameter("constant", 1.0);
 
                    Filler<T> filler = Filler<T>.Create(m_cuda, m_log, fp);
                    filler.Fill(blobScale);
                }
                m_colBlobs.Add(blobScale);
            }
 
            if (p.bias_term)
            {
                LayerParameter pb = new LayerParameter(LayerParameter.LayerType.BIAS);
                pb.bias_param.axis = p.axis;
                pb.bias_param.num_axes = (colBottom.Count > 1) ? colBottom[1].num_axes : p.num_axes;
                pb.bias_param.filler = p.bias_filler;
 
                m_colBiasBottomVec = new BlobCollection<T>();
                m_colBiasBottomVec.Add(colBottom[0]);
 
                m_biasLayer = new BiasLayer<T>(m_cuda, m_log, pb);
                m_biasLayer.Setup(m_colBiasBottomVec, colTop);
 
                shareLayerBlobs(m_biasLayer);
 
                m_nBiasParamId = m_colBlobs.Count;
                m_colBlobs.Add(m_biasLayer.blobs[0]);
                m_rgbBiasPropagateDown = Utility.Create<bool>(1, false);
            }
 
            m_rgbParamPropagateDown = new DictionaryMap<bool>(m_colBlobs.Count(), true);
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            ScaleParameter p = m_param.scale_param;
            Blob<T> blobScale = (colBottom.Count > 1) ? colBottom[1] : m_colBlobs[0];
 
            // Always set axis == 0 in special case where bias is a scalar
            // (num_axes == 0.  Mathematically eqiuvalent for any choice of axis, so the
            // actual setting can be safely ignored; and computation is most efficient
            // with axis == 0 and (therefore) outer_dim == 1. (Setting m_nAxis to
            // bottom[0].num_axes - 1, giving inner_dim_ == 1, would be equally
            // performant.)
            m_nAxis = (blobScale.num_axes == 0) ? 0 : colBottom[0].CanonicalAxisIndex(p.axis);
            m_log.CHECK_GE(colBottom[0].num_axes, m_nAxis + blobScale.num_true_axes, "scale blob's shape extends past bottom[0]'s shape when applied starting with bottom[0] axis = " + m_nAxis.ToString());
 
            for (int i = 0; i < blobScale.num_true_axes; i++)
            {
                m_log.CHECK_EQ(colBottom[0].shape(m_nAxis + i), blobScale.shape(i), "dimension mismatch between bottom[0]->shape(" + (m_nAxis + i).ToString() + ") and scale->shape(" + i.ToString() + ")");
            }
 
            m_nOuterDim = colBottom[0].count(0, m_nAxis);
            m_nScaleDim = blobScale.count();
            m_nInnerDim = colBottom[0].count(m_nAxis + blobScale.num_true_axes);
 
            if (colBottom[0] == colTop[0])  // in-place computation
                m_blobTemp.ReshapeLike(colBottom[0]);
            else
                colTop[0].ReshapeLike(colBottom[0]);
 
            m_blobSumResult.Reshape(new List<int>() { m_nOuterDim * m_nScaleDim });
            int nSumMultSize = Math.Max(m_nOuterDim, m_nInnerDim);
            m_blobSumMultiplier.Reshape(new List<int>() { nSumMultSize });
            m_blobSumMultiplier.SetData(1.0);
 
            if (m_biasLayer != null)
            {
                m_colBiasBottomVec[0] = colTop[0];
                m_biasLayer.Reshape(m_colBiasBottomVec, colTop);
            }
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            if (colBottom[0] == colTop[0])
            {
                // in-place computation; need to store bottom data before overwriting it.
                // Note that this is only necessary for backward; we could skip this if not
                // doing backward, but Caffe currently provides no way of knowing whether
                // we'll need to do backward at the time of the forward call.
                m_blobTemp.CopyFrom(colBottom[0]);
            }
 
            long hScaleData = (colBottom.Count > 1) ? colBottom[1].gpu_data : m_colBlobs[0].gpu_data;
            long hTopData = colTop[0].mutable_gpu_data;
            int nCount = colTop[0].count();
            long hBottomData = colBottom[0].gpu_data;
 
            if (m_biasLayer != null)
            {
                long hBiasData = m_colBlobs[m_nBiasParamId].gpu_data;
                m_cuda.scale_fwd(nCount, hBottomData, hScaleData, m_nScaleDim, m_nInnerDim, hTopData, hBiasData);
            }
            else
            {
                m_cuda.scale_fwd(nCount, hBottomData, hScaleData, m_nScaleDim, m_nInnerDim, hTopData);
            }
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            if (m_biasLayer != null && m_rgbParamPropagateDown[m_rgbParamPropagateDown.Count - 1])
                m_biasLayer.Backward(colTop, m_rgbBiasPropagateDown, m_colBiasBottomVec);
 
            bool bScaleParam = (colBottom.Count == 1) ? true : false;
            Blob<T> blobScale = (bScaleParam) ? m_colBlobs[0] : colBottom[1];
 
            if ((!bScaleParam && rgbPropagateDown[1]) || (bScaleParam && m_rgbParamPropagateDown[0]))
            {
                long hTopDiff = colTop[0].gpu_diff;
                bool bInPlace = (colBottom[0] == colTop[0]) ? true : false;
                long hBottomData = (bInPlace) ? m_blobTemp.gpu_data : colBottom[0].gpu_data;
 
                // Hack: store big eltwise product in bottom[0].diff, except in the special
                // case where this layer itself does the eltwise product, in which case we
                // can store it directly in the scale diff, and we're done.
                // If we're computing in-place (and not doing eltwise computation), this
                // hack doesn't work and we store the product in temp_.
                bool bIsEltwise = (colBottom[0].count() == blobScale.count()) ? true : false;
                long hProduct = (bIsEltwise) ? blobScale.mutable_gpu_diff : ((bInPlace) ? m_blobTemp.mutable_gpu_data : colBottom[0].mutable_gpu_diff);
                long hSumMult = m_blobSumMultiplier.gpu_data;
 
                m_cuda.mul(colTop[0].count(), hTopDiff, hBottomData, hProduct);
 
                if (!bIsEltwise)
                {
                    long hSumResult = 0;
 
                    if (m_nInnerDim == 1)
                    {
                        hSumResult = hProduct;
                    }
                    else if (m_blobSumResult.count() == 1)
                    {
                        double dfScaleDiff = convertD(blobScale.GetDiff(0));
 
                        if (bScaleParam)
                        {
                            T fDot = m_cuda.dot(m_nInnerDim, hProduct, hSumMult);
                            dfScaleDiff += convertD(fDot);
                            blobScale.SetDiff(dfScaleDiff, 0);
                        }
                        else
                        {
                            T fDot = m_cuda.dot(m_nInnerDim, hProduct, hSumMult);
                            blobScale.SetDiff(convertD(fDot), 0);
                        }
                    }
                    else
                    {
                        hSumResult = (m_nOuterDim == 1) ? blobScale.mutable_gpu_diff : m_blobSumResult.mutable_gpu_data;
                        m_cuda.gemv(false, m_blobSumResult.count(), m_nInnerDim, m_tOne, hProduct, hSumMult, m_tZero, hSumResult);
                    }
 
                    if (m_nOuterDim != 1)
                    {
                        if (m_nScaleDim == 1)
                        {
                            double dfScaleDiff = convertD(blobScale.GetDiff(0));
 
                            if (bScaleParam)
                            {
                                T fDot = m_cuda.dot(m_nOuterDim, hSumMult, hSumResult);
                                dfScaleDiff += convertD(fDot);
                                blobScale.SetDiff(dfScaleDiff, 0);
                            }
                            else
                            {
                                T fDot = m_cuda.dot(m_nOuterDim, hSumMult, hSumResult);
                                blobScale.SetDiff(convertD(fDot), 0);
                            }
                        }
                        else
                        {
                            long hScaleDiff = blobScale.mutable_gpu_diff;
                            m_cuda.gemv(true, m_nOuterDim, m_nScaleDim, m_tOne, hSumResult, hSumMult, (bScaleParam) ? m_tOne : m_tZero, hScaleDiff);
                        }
                    }
                }
            }
 
            if (rgbPropagateDown[0])
            {
                int nCount = colTop[0].count();
                long hTopDiff = colTop[0].gpu_diff;
                long hScaleData = blobScale.gpu_data;
                long hBottomDiff = colBottom[0].mutable_gpu_diff;
 
                m_cuda.scale_fwd(nCount, hTopDiff, hScaleData, m_nScaleDim, m_nInnerDim, hBottomDiff);
            }
        }
    }
}