ReshapeTemporalLayer.cs

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Net.Http.Headers;
using System.Reflection;
using System.Text;
using MyCaffe.basecode;
using MyCaffe.common;
using MyCaffe.param;
using MyCaffe.param.tft;
 
namespace MyCaffe.layers.tft
{
    public class ReshapeTemporalLayer<T> : Layer<T>
    {
        ReshapeTemporalParameter.MODE m_mode;
        int m_nNumSamples;
        int m_nNumRepeatCount;
        int m_nForcedRepeatCount;
        List<int> m_rgShape = new List<int>(4);
        Dictionary<string, List<int>> m_rgShapes = new Dictionary<string, List<int>>();
        Blob<T> m_blobTimeDistributedContext;
        Blob<T> m_blobWork;
 
        public ReshapeTemporalLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.RESHAPE_TEMPORAL;
        }
 
        protected override void dispose()
        {
            dispose(ref m_blobTimeDistributedContext);
            dispose(ref m_blobWork);
        }
 
        protected override void setup_internal_blobs(BlobCollection<T> col)
        {
            if (col.Count > 0)
                return;
 
            if (m_blobTimeDistributedContext != null)
                col.Add(m_blobTimeDistributedContext);
        }
 
        public override int MinBottomBlobs
        {
            get { return 1; }
        }
 
        public override int MaxBottomBlobs
        {
            get { return 2; }
        }
 
        public override int MinTopBlobs
        {
            get { return 1; }
        }
 
        public override int MaxTopBlobs
        {
            get { return 3; }
        }
 
        private void replicate_along_time_fwd(Blob<T> bBtm, Blob<T> bTop, int nTimeSteps, bool bTemporalRepeat, bool bReshapeOnly = false)
        {
            m_rgShape.Clear();
 
            if (bTemporalRepeat)
            {
                m_rgShape.Add(m_nNumSamples);
                m_rgShape.Add(nTimeSteps);
                m_rgShape.Add(bBtm.shape(1));
                bTop.Reshape(m_rgShape);
 
                if (!bReshapeOnly)
                {
                    int nInnerNum = bBtm.count(1);
                    for (int i = 0; i < nTimeSteps; i++)
                    {
                        m_cuda.channel_copy(bBtm.count(), m_nNumSamples, 1, nTimeSteps, nInnerNum, i, bTop.mutable_gpu_data, bBtm.gpu_data, DIR.BWD);
                    }
                }
            }
            else
            {
                m_rgShape.Add(nTimeSteps);
                m_rgShape.Add(m_nNumSamples);
                m_rgShape.Add(bBtm.shape(1));
                bTop.Reshape(m_rgShape);
 
                if (!bReshapeOnly)
                {
                    int nInnerNum = bBtm.count(1);
                    for (int i = 0; i < nTimeSteps; i++)
                    {
                        m_cuda.channel_copy(bBtm.count(), 1, 1, nTimeSteps, m_nNumSamples * nInnerNum, i, bTop.mutable_gpu_data, bBtm.gpu_data, DIR.BWD);
                    }
                }
            }
        }
 
        private void replicate_along_time_bwd(Blob<T> bBtm, Blob<T> bTop, int nTimeSteps, bool bTemporalRepeat)
        {
            int nInnerNum = bBtm.count(1);
 
            if (bTemporalRepeat)    
                m_cuda.channel_sum(bTop.count(), m_nNumSamples, nTimeSteps, nInnerNum, bTop.gpu_diff, bBtm.mutable_gpu_diff, true);
            else
                m_cuda.channel_sum(bTop.count(), 1, nTimeSteps, m_nNumSamples * nInnerNum, bTop.gpu_diff, bBtm.mutable_gpu_diff, true);
 
            //m_cuda.channel_copy(bBtm.count(), m_nNumSamples, 1, nTimeSteps, nInnerNum, 0, bTop.mutable_gpu_diff, bBtm.gpu_diff, DIR.FWD);
 
            //for (int i = 1; i < nTimeSteps; i++)
            //{
            //    m_cuda.channel_add(bBtm.count(), m_nNumSamples, 1, nTimeSteps, nInnerNum, i, bTop.mutable_gpu_diff, bBtm.mutable_gpu_diff, DIR.FWD);
            //}
        }
 
        private void stack_time_steps_along_batch_fwd(Blob<T> bBtm, Blob<T> bTop, bool bResizeOnly = false)
        {
            if (!bResizeOnly)
                bTop.CopyFrom(bBtm, false, true);
 
            m_rgShape.Clear();
            m_rgShape.Add(bBtm.shape(0) * bBtm.shape(1));
            m_rgShape.Add(bBtm.count(2));
            bTop.Reshape(m_rgShape);
        }
 
        private void stack_time_steps_along_batch_bwd(Blob<T> bBtm, Blob<T> bTop)
        {
            bBtm.CopyFrom(bTop, true, false, 0, true);
        }
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            m_mode = m_param.reshape_temporal_param.mode;
 
            if (m_mode == param.tft.ReshapeTemporalParameter.MODE.BEFORE)
            {
                m_nNumSamples = colBottom[0].num;
                m_nForcedRepeatCount = m_param.reshape_temporal_param.forced_repeat_count;
 
                if (m_nForcedRepeatCount >= 0)
                    m_nNumRepeatCount = m_nForcedRepeatCount;
                else
                    m_nNumRepeatCount = colBottom[0].shape(1);
 
                // replicate the selection signal along time
                if (colBottom.Count > 1)
                {
                    m_blobWork = new Blob<T>(m_cuda, m_log);
                    m_blobWork.Name = "work";
 
                    if (m_nNumRepeatCount > 0)
                    {
                        m_blobTimeDistributedContext = new Blob<T>(m_cuda, m_log);
                        m_blobTimeDistributedContext.Name = m_param.name + ".tdctx";
                        replicate_along_time_fwd(colBottom[1], m_blobTimeDistributedContext, m_nNumRepeatCount, m_nForcedRepeatCount < 0, true);
                        stack_time_steps_along_batch_fwd(m_blobTimeDistributedContext, colTop[1], true);
                    }
                    else
                    {
                        stack_time_steps_along_batch_fwd(colBottom[1], colTop[1], true);
                    }
 
                    colTop[1].SetParameter("num_samples", m_nNumSamples);
                    colTop[1].SetParameter("num_temporal_steps", m_nNumRepeatCount);
                    colTop[1].SetParameter("forced_temporal_steps", m_nForcedRepeatCount);
                }
 
                // Apply the same selection module on all timesteps by stacking the time dimension with the batch dimension
                stack_time_steps_along_batch_fwd(colBottom[0], colTop[0], true);
                colTop[0].SetParameter("num_samples", m_nNumSamples);
                colTop[0].SetParameter("num_temporal_steps", colBottom[0].shape(1));
            }
            else
            {
                m_nNumSamples = (int)colBottom[0].GetParameter("num_samples");
                int nTemporalSteps = (int)colBottom[0].GetParameter("num_temporal_steps");
 
                int nCount = colBottom[0].count();
                int nDim = m_nNumSamples * nTemporalSteps;
                m_rgShape.Clear();
                m_rgShape.Add(m_nNumSamples);
                m_rgShape.Add(nTemporalSteps);
                m_rgShape.Add(nCount / nDim);
                colTop[0].Reshape(m_rgShape);
 
                int nIdx = 1;
                if (m_param.reshape_temporal_param.enable_clip_output)
                {
                    m_log.CHECK_GT(colTop.Count, nIdx, "There must be at least " + (nIdx + 1).ToString() + " tops for the enable clip output!");
                    m_rgShape.Clear();
                    m_rgShape.Add(m_nNumSamples);
                    m_rgShape.Add(nTemporalSteps);
                    colTop[nIdx].Reshape(m_rgShape);
                    nIdx++;
                }
 
                if (colBottom.Count > 1)
                {
                    m_nNumRepeatCount = (int)colBottom[1].GetParameter("num_temporal_steps");
                    m_nForcedRepeatCount = (int)colBottom[1].GetParameter("forced_temporal_steps");
 
                    if (m_param.reshape_temporal_param.enable_weight_output)
                    {
                        m_log.CHECK_GT(colTop.Count, nIdx, "There must be at least " + (nIdx + 1).ToString() + " tops for the enable clip output!");
                        nCount = colBottom[1].count();
                        m_rgShape.Clear();
 
                        if (m_nForcedRepeatCount >= 0)
                        {
                            m_rgShape.Add(m_nNumSamples);
                            m_rgShape.Add(nCount / nDim);
                        }
                        else
                        {
                            m_rgShape.Add(m_nNumSamples);
                            if (m_nNumRepeatCount > 0)
                                m_rgShape.Add(m_nNumRepeatCount);
                            m_rgShape.Add(nCount / nDim);
                        }
 
                        colTop[nIdx].Reshape(m_rgShape);
                    }
                }
            }
 
            setup_internal_blobs(m_colInternalBlobs);
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            if (m_mode == ReshapeTemporalParameter.MODE.BEFORE)
            {
                // replicate the selection signal along time
                if (colBottom.Count > 1)
                {
                    if (m_nNumRepeatCount > 0)
                    {
                        replicate_along_time_fwd(colBottom[1], m_blobTimeDistributedContext, m_nNumRepeatCount, m_nForcedRepeatCount < 0, true);
                        m_blobWork.ReshapeLike(colBottom[1]);
                        stack_time_steps_along_batch_fwd(m_blobTimeDistributedContext, colTop[1], true);
                    }
                    else
                    {
                        stack_time_steps_along_batch_fwd(colBottom[1], colTop[1], true);
                    }
                }
 
                // Apply the same selection module on all timesteps by stacking the time dimension with the batch dimension
                stack_time_steps_along_batch_fwd(colBottom[0], colTop[0], true);
            }
            else
            {
                int nTemporalSteps = (int)colBottom[0].GetParameter("num_temporal_steps");
                int nCount = colBottom[0].count();
                int nDim = m_nNumSamples * nTemporalSteps;
                m_rgShape.Clear();
                m_rgShape.Add(m_nNumSamples);
                m_rgShape.Add(nTemporalSteps);
                m_rgShape.Add(nCount / nDim);
                colTop[0].Reshape(m_rgShape);
 
                int nIdx = 1;
                if (m_param.reshape_temporal_param.enable_clip_output)
                {
                    m_log.CHECK_GT(colTop.Count, nIdx, "There must be at least " + (nIdx + 1).ToString() + " tops for the enable clip output!");
                    m_rgShape.Clear();
                    m_rgShape.Add(m_nNumSamples);
                    m_rgShape.Add(nTemporalSteps);
                    colTop[nIdx].Reshape(m_rgShape);
                    nIdx++;
                }
 
                if (colBottom.Count > 1)
                {
                    m_nNumRepeatCount = (int)colBottom[1].GetParameter("num_temporal_steps");
                    m_nForcedRepeatCount = (int)colBottom[1].GetParameter("forced_temporal_steps");
 
                    if (m_param.reshape_temporal_param.enable_weight_output)
                    {
                        m_log.CHECK_GT(colTop.Count, nIdx, "There must be at least " + (nIdx + 1).ToString() + " tops for the enable clip output!");
                        nCount = colBottom[1].count();
                        m_rgShape.Clear();
 
                        if (m_nForcedRepeatCount >= 0)
                        {
                            m_rgShape.Add(m_nNumSamples);
                            m_rgShape.Add(nCount / nDim);
                        }
                        else
                        {
                            m_rgShape.Add(m_nNumSamples);
                            if (m_nNumRepeatCount > 0)
                                m_rgShape.Add(m_nNumRepeatCount);
                            m_rgShape.Add(nCount / nDim);
                        }
 
                        colTop[nIdx].Reshape(m_rgShape);
                    }
                }
            }
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            if (m_mode == ReshapeTemporalParameter.MODE.BEFORE)
            {
                if (colBottom.Count > 1)
                {
                    // replicate the selection signal along time
                    if (m_nNumRepeatCount > 0)
                    {
                        replicate_along_time_fwd(colBottom[1], m_blobTimeDistributedContext, m_nNumRepeatCount, m_nForcedRepeatCount < 0);
                        stack_time_steps_along_batch_fwd(m_blobTimeDistributedContext, colTop[1]);
                    }
                    else
                    {
                        stack_time_steps_along_batch_fwd(colBottom[1], colTop[1]);
                    }
                }
 
                // Apply the same selection module on all timesteps by stacking the time dimension with the batch dimension
                stack_time_steps_along_batch_fwd(colBottom[0], colTop[0]);
            }
            else
            {
                colTop[0].CopyFrom(colBottom[0], false, false, 0, true);
 
                int nIdx = 1;
                if (m_param.reshape_temporal_param.enable_clip_output)
                {
                    colTop[nIdx].SetData(1);
                    nIdx++;
                }
 
                if (m_param.reshape_temporal_param.enable_weight_output)
                {
                    colTop[nIdx].CopyFrom(colBottom[1], false, false, 0, true);
                }
            }
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            if (m_mode == ReshapeTemporalParameter.MODE.BEFORE)
            {
                // Apply the same selection module on all timesteps by stacking the time dimension with the batch dimension
                stack_time_steps_along_batch_bwd(colBottom[0], colTop[0]);
 
                // replicate the static selection signal along time
                if (colBottom.Count > 1)
                {
                    if (m_nNumRepeatCount > 0)
                    {
                        stack_time_steps_along_batch_bwd(m_blobTimeDistributedContext, colTop[1]);
                        replicate_along_time_bwd(colBottom[1], m_blobTimeDistributedContext, m_nNumRepeatCount, m_nForcedRepeatCount < 0);
                    }
                    else
                    {
                        stack_time_steps_along_batch_bwd(colBottom[1], colTop[1]);
                    }
                }
            }
            else
            {
                colBottom[0].CopyFrom(colTop[0], true, false, 0, true);
                int nIdx = 1;
 
                if (m_param.reshape_temporal_param.enable_clip_output)
                    nIdx++;
 
                if (colBottom.Count > 1)
                    colBottom[1].CopyFrom(colTop[nIdx], true, false, 0, true);
            }
        }
    }
}