MergeLayer.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using MyCaffe.basecode;
using MyCaffe.common;
using MyCaffe.param;
using MyCaffe.param.beta;
 
namespace MyCaffe.layers.beta
{
    public class MergeLayer<T> : Layer<T>
    {
        int m_nCopyAxis = 0;
        int m_nOrderMajorAxis = 0;
        int m_nCopyCount = 1;
        int m_nSrcStartIdx1 = 0;
        int m_nDstStartIdx1 = 0;
        int m_nCopyDim1 = 1;
        int m_nSrcStartIdx2 = 0;
        int m_nDstStartIdx2 = 0;
        int m_nCopyDim2 = 1;
        int m_nSrcSpatialDim1 = 1;
        int m_nSrcSpatialDim2 = 1;
        int m_nSrcSpatialDimStartIdx1 = 0;
        int m_nDstSpatialDimStartIdx1 = 0;
        int m_nSrcSpatialDimStartIdx2 = 0;
        int m_nDstSpatialDimStartIdx2 = 0;
        int m_nSpatialDimCopyCount = -1;
        int m_nDstSpatialDim = 1;
 
        public MergeLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.MERGE;
        }
 
        protected override void dispose()
        {
            base.dispose();
        }
 
        public override int ExactNumBottomBlobs
        {
            get { return 2; }
        }
 
        public override int ExactNumTopBlobs
        {
            get { return 1; }
        }
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            m_nCopyAxis = m_param.merge_param.copy_axis;
            m_nOrderMajorAxis = m_param.merge_param.order_major_axis;
            m_nCopyCount = m_param.merge_param.copy_count;
            m_nSrcStartIdx1 = Utility.CanonicalAxisIndex(m_param.merge_param.src_start_idx1, colBottom[0].shape()[m_nCopyAxis]);
            m_nDstStartIdx1 = Utility.CanonicalAxisIndex(m_param.merge_param.dst_start_idx1, colBottom[0].shape()[m_nCopyAxis]);
            m_nCopyDim1 = m_param.merge_param.copy_dim1;
            m_nSrcStartIdx2 = Utility.CanonicalAxisIndex(m_param.merge_param.src_start_idx2, colBottom[1].shape()[m_nCopyAxis]);
            m_nDstStartIdx2 = Utility.CanonicalAxisIndex(m_param.merge_param.dst_start_idx2, colBottom[1].shape()[m_nCopyAxis]);
            m_nCopyDim2 = m_param.merge_param.copy_dim2;
 
            m_nSrcSpatialDim1 = Utility.GetSpatialDim(colBottom[0].shape(), Math.Max(m_nCopyAxis, m_nOrderMajorAxis) + 1);
            m_nSrcSpatialDim2 = Utility.GetSpatialDim(colBottom[1].shape(), Math.Max(m_nCopyAxis, m_nOrderMajorAxis) + 1);
            m_nDstSpatialDim = m_param.merge_param.dst_spatialdim;
            if (m_nDstSpatialDim <= 0)
                m_nDstSpatialDim = 1;
 
            m_nSrcSpatialDimStartIdx1 = Utility.CanonicalAxisIndex(m_param.merge_param.src_spatialdim_start_idx1, m_nSrcSpatialDim1);
            m_nSrcSpatialDimStartIdx2 = Utility.CanonicalAxisIndex(m_param.merge_param.src_spatialdim_start_idx1, m_nSrcSpatialDim2);
            m_nDstSpatialDimStartIdx1 = Utility.CanonicalAxisIndex(m_param.merge_param.dst_spatialdim_start_idx1, m_nDstSpatialDim);
            m_nDstSpatialDimStartIdx2 = Utility.CanonicalAxisIndex(m_param.merge_param.dst_spatialdim_start_idx1, m_nDstSpatialDim);
            m_nSpatialDimCopyCount = m_param.merge_param.spatialdim_copy_count;
 
            m_log.CHECK_GT(m_nCopyCount, 0, "The copy_count must be > 0!");
 
            m_log.CHECK_GT(m_nCopyDim1, 0, "The copy dim must be > 0!");
            m_log.CHECK_GT(m_nCopyDim2, 0, "The copy dim must be > 0!");
            m_log.CHECK_GE(m_nSrcStartIdx1, 0, "The start_idx1 must be >= 0!");
            m_log.CHECK_GE(m_nSrcStartIdx2, 0, "The start_idx1 must be >= 0!");
            m_log.CHECK_GE(m_nDstStartIdx1, 0, "The start_idx1 must be >= 0!");
            m_log.CHECK_GE(m_nDstStartIdx2, 0, "The start_idx1 must be >= 0!");
 
            m_log.CHECK_GE(m_nCopyAxis, 0, "The copy axis must be >= 0!");
            m_log.CHECK_LE(m_nCopyAxis, 1, "The copy axis must be <= 1!");
            m_log.CHECK_EQ(m_nOrderMajorAxis, 1, "Currently, only order_major_axis=1 is supported.");
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            m_log.CHECK_EQ(colBottom.Count, 2, "Two inputs are expected!");
            List<int> rgTopShape = MergeParameter.Reshape(m_log, m_param.merge_param, colBottom[0].shape(), colBottom[1].shape());
 
            colTop[0].Reshape(rgTopShape);
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            long hBtmData0 = colBottom[0].gpu_data;
            long hBtmData1 = colBottom[1].gpu_data;
            long hTopData = colTop[0].mutable_gpu_data;
 
            int nDstStep = colTop[0].shape()[m_nCopyAxis];
            int nSrcStep1 = colBottom[0].shape()[m_nCopyAxis];
            m_cuda.copy_sequence(colBottom[0].count(), hBtmData0, nSrcStep1, m_nSrcStartIdx1, m_nCopyCount, m_nCopyDim1, hTopData, nDstStep, m_nDstStartIdx1, m_nSrcSpatialDim1, m_nDstSpatialDim, m_nSrcSpatialDimStartIdx1, m_nDstSpatialDimStartIdx1, m_nSpatialDimCopyCount);
 
            int nSrcStep2 = colBottom[1].shape()[m_nCopyAxis];
            m_cuda.copy_sequence(colBottom[1].count(), hBtmData1, nSrcStep2, m_nSrcStartIdx2, m_nCopyCount, m_nCopyDim2, hTopData, nDstStep, m_nDstStartIdx2, m_nSrcSpatialDim2, m_nDstSpatialDim, m_nSrcSpatialDimStartIdx2, m_nDstSpatialDimStartIdx2, m_nSpatialDimCopyCount);
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            colBottom[0].SetDiff(0);
            colBottom[1].SetDiff(0);
            long hBtmDiff0 = colBottom[0].mutable_gpu_diff;
            long hBtmDiff1 = colBottom[1].mutable_gpu_diff;
            long hTopData = colTop[0].gpu_diff;
 
            int nSrcStep = colTop[0].shape()[m_nCopyAxis];
            int nDstStep1 = colBottom[0].shape()[m_nCopyAxis];
            m_cuda.copy_sequence(colTop[0].count(), hTopData, nSrcStep, m_nDstStartIdx1, m_nCopyCount, m_nCopyDim1, hBtmDiff0, nDstStep1, m_nSrcStartIdx1, m_nDstSpatialDim, m_nSrcSpatialDim1, m_nDstSpatialDimStartIdx1, m_nSrcSpatialDimStartIdx1, m_nSpatialDimCopyCount);
 
            int nDstStep2 = colBottom[1].shape()[m_nCopyAxis];
            m_cuda.copy_sequence(colTop[0].count(), hTopData, nSrcStep, m_nDstStartIdx2, m_nCopyCount, m_nCopyDim2, hBtmDiff1, nDstStep2, m_nSrcStartIdx2, m_nDstSpatialDim, m_nSrcSpatialDim2, m_nDstSpatialDimStartIdx2, m_nSrcSpatialDimStartIdx2, m_nSpatialDimCopyCount);
        }
    }
}