ConcatLayer.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 ConcatLayer<T> : Layer<T>
    {
        int m_nNumConcats;
        int m_nConcatInputSize;
        int m_nConcatAxis;
 
        public ConcatLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.CONCAT;
        }
 
        public override int MinBottomBlobs
        {
            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)
        {
            int nNumAxes = colBottom[0].num_axes;
 
            if (m_param.concat_param.concat_dim.HasValue)
            {
                m_nConcatAxis = (int)m_param.concat_param.concat_dim.Value;
                // Don't allow negative indexing for concat_dim, a uint -- almost certainly
                // unintended.
                m_log.CHECK_GE(m_nConcatAxis, 0, "Casting concat_dim from uint to int produced a negative result; concat_dim must be > 0.");
                m_log.CHECK_LT(m_nConcatAxis, nNumAxes, "concat_dim out of range.");
            }
            else
            {
                m_nConcatAxis = colBottom[0].CanonicalAxisIndex(m_param.concat_param.axis);
            }
 
            // Initialize with the first blob.
            List<int> rgTopShape = Utility.Clone<int>(colBottom[0].shape());
            m_nNumConcats = colBottom[0].count(0, m_nConcatAxis);
            m_nConcatInputSize = colBottom[0].count(m_nConcatAxis + 1);
 
            int nBottomCountSum = colBottom[0].count();
 
            for (int i = 1; i < colBottom.Count; i++)
            {
                m_log.CHECK_EQ(nNumAxes, colBottom[i].num_axes, "All inputs must have the same # axes.");
 
                for (int j = 0; j < nNumAxes; j++)
                {
                    if (j == m_nConcatAxis)
                        continue;
 
                    m_log.CHECK_EQ(rgTopShape[j], colBottom[i].shape(j), "All inputs must have the same shape, except at concat_axis.  You might try switching between the ONNX(p) and CAFFE(t) type pooling sizing methods.");
                }
 
                nBottomCountSum += colBottom[i].count();
                rgTopShape[m_nConcatAxis] += colBottom[i].shape(m_nConcatAxis);
            }
 
            colTop[0].Reshape(rgTopShape);
            m_log.CHECK_EQ(nBottomCountSum, colTop[0].count(), "The bottomCountSums should equal the top[0].count.");
 
            if (colBottom.Count == 1)
            {
                colTop[0].ShareData(colBottom[0]);
                colTop[0].ShareDiff(colBottom[0]);
            }
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            if (colBottom.Count == 1)
                return;
 
            long hTopData = colTop[0].mutable_gpu_data;
            int nOffsetConcatAxis = 0;
            int nTopConcatAxis = colTop[0].shape(m_nConcatAxis);
 
            for (int i = 0; i < colBottom.Count; i++)
            {
                long hBottomData = colBottom[i].gpu_data;
                int nBottomConcatAxis = colBottom[i].shape(m_nConcatAxis);
                int nBottomConcatSize = nBottomConcatAxis * m_nConcatInputSize;
                int nCount = nBottomConcatSize * m_nNumConcats;
 
                m_cuda.concat_fwd(nCount, hBottomData, m_nNumConcats, m_nConcatInputSize, nTopConcatAxis, nBottomConcatAxis, nOffsetConcatAxis, hTopData);
                nOffsetConcatAxis += nBottomConcatAxis;
            }
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            if (colBottom.Count == 1)
                return;
 
            long hTopDiff = colTop[0].gpu_diff;
            int nOffsetConcatAxis = 0;
            int nTopConcatAxis = colTop[0].shape(m_nConcatAxis);
 
            for (int i = 0; i < colBottom.Count; i++)
            {
                int nBottomConcatAxis = colBottom[i].shape(m_nConcatAxis);
 
                if (rgbPropagateDown[i])
                {
                    long hBottomDiff = colBottom[i].mutable_gpu_diff;
                    int nBottomConcatSize = nBottomConcatAxis * m_nConcatInputSize;
                    int nCount = nBottomConcatSize * m_nNumConcats;
 
                    m_cuda.concat_bwd(nCount, hTopDiff, m_nNumConcats, m_nConcatInputSize, nTopConcatAxis, nBottomConcatAxis, nOffsetConcatAxis, hBottomDiff);
                }
 
                nOffsetConcatAxis += nBottomConcatAxis;
            }
        }
    }
}