SplitLayer.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 SplitLayer<T> : Layer<T>
    {
        int m_nCount;
 
        public SplitLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.SPLIT;
        }
 
        public override int ExactNumBottomBlobs
        {
            get { return 1; }
        }
 
        public override int MinTopBlobs
        {
            get { return 1; }
        }
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            m_bConvertBottom = false;
 
            // Copy data during setup through
            // forward pass to ensure variables
            // are passed through.
            Reshape(colBottom, colTop);
            forward(colBottom, colTop);
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            m_nCount = colBottom[0].count();
 
            for (int i = 0; i < colTop.Count; i++)
            {
                // Do not allow in-place computation in the SplitLayer.  Instead, share data by 
                // reference in the forward pass, and keep separeate diff allocations in
                // the backward pass.  (Technically, it should be possible to share the diff
                // blob of the first split output with the input, but this seems to cause
                // some strange effects in practice...)
                m_log.CHECK(colTop[i].gpu_data != colBottom[0].gpu_data, "Layer does not allow in-place computation.");
                colTop[i].ReshapeLike(colBottom[0], colBottom[0].HalfSize);
                m_log.CHECK_EQ(m_nCount, colTop[i].count(), "The count should equal the top[i].count().");
            }
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            int nCount = colBottom[0].count();
            long hBottom = colBottom[0].gpu_data;
 
            for (int i = 0; i < colTop.Count; i++)
            {
                m_cuda.copy(nCount, hBottom, colTop[i].mutable_gpu_data);
            }
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            if (!rgbPropagateDown[0])
                return;
 
            if (colTop.Count == 1)
            {
                m_cuda.copy(m_nCount, colTop[0].gpu_diff, colBottom[0].mutable_gpu_diff);
                return;
            }
 
            m_cuda.add(m_nCount, colTop[0].gpu_diff, colTop[1].gpu_diff, colBottom[0].mutable_gpu_diff);
 
            // Add remaining top blob diffs.
            for (int i = 2; i < colTop.Count; i++)
            {
                long hTopDiff = colTop[i].gpu_diff;
                long hBottomDiff = colBottom[0].mutable_gpu_diff;
 
                m_cuda.axpy(m_nCount, m_tOne, hTopDiff, hBottomDiff);
            }
        }
    }
}