ChannelEmbeddingLayer.cs

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Text;
using MyCaffe.basecode;
using MyCaffe.common;
using MyCaffe.param;
 
namespace MyCaffe.layers.tft
{
    public class ChannelEmbeddingLayer<T> : Layer<T>
    {
        Layer<T> m_numericLayer = null;
        Layer<T> m_categoricalLayer = null;
        BlobCollection<T> m_colBtm = new BlobCollection<T>();
        BlobCollection<T> m_colNumericTop = new BlobCollection<T>();
        BlobCollection<T> m_colCategoricalTop = new BlobCollection<T>();
        List<int> m_rgShape = new List<int>(4);
 
        public ChannelEmbeddingLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.CHANNEL_EMBEDDING;
        }
 
        protected override void dispose()
        {
            dispose(ref m_numericLayer);
            dispose(ref m_categoricalLayer);
 
            m_colNumericTop.Dispose();
            m_colCategoricalTop.Dispose();
        }
 
        protected override void setup_internal_blobs(BlobCollection<T> col)
        {
            if (col.Count > 0)
                return;
        }
 
        public override int MinBottomBlobs
        {
            get { return 1; }
        }
 
        public override int MaxBottomBlobs
        {
            get { return 2; }
        }
 
        public override int ExactNumTopBlobs
        {
            get { return 1; }
        }
 
        private void getBlobs(BlobCollection<T> colBottom, out Blob<T> blobNumeric, out Blob<T> blobCategorical)
        {
            blobNumeric = null;
            blobCategorical = null;
 
            m_log.CHECK_EQ(colBottom.Count, 2, "The bottom must have a count = 2.");
 
            if (m_param.numeric_trans_param.num_input > 0)
            {
                m_log.CHECK_GT(colBottom[0].count(), 0, "The bottom(0) must have a count > 0!");
                blobNumeric = colBottom[0];
            }
            if (m_param.categorical_trans_param.num_input > 0)
            {
                m_log.CHECK_GT(colBottom[1].count(), 0, "The bottom(1) must have a count > 0!");
                blobCategorical = colBottom[1];
            }
            
            if (blobNumeric == null && blobCategorical == null)
                m_log.FAIL("At least one of the numeric or categorical num_input must be > 0.");
        }
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            Blob<T> blobNumeric = null;
            Blob<T> blobCategorical = null;
 
            m_log.CHECK_EQ(m_param.numeric_trans_param.state_size, m_param.categorical_trans_param.state_size, "The numeric and categorical parameters must have the same sate_size.");
 
            getBlobs(colBottom, out blobNumeric, out blobCategorical);
 
            if (blobNumeric != null)
            {
                m_colNumericTop.Clear();
                for (int i=0; i<m_param.numeric_trans_param.num_input; i++)
                {
                    Blob<T> blobTop = new Blob<T>(m_cuda, m_log);
                    m_colNumericTop.Add(blobTop);
                }
 
                LayerParameter p = new LayerParameter(LayerParameter.LayerType.NUMERIC_TRANS, m_param.name + ".numeric");
                p.numeric_trans_param.Copy(m_param.numeric_trans_param);
                m_numericLayer = Layer<T>.Create(m_cuda, m_log, convertLayerParam(p, m_param), null) as NumericTransformationLayer<T>;
 
                m_colBtm.Clear();
                m_colBtm.Add(blobNumeric);
                m_numericLayer.LayerSetUp(m_colBtm, m_colNumericTop);
                blobs.Add(m_numericLayer.blobs);
            }
 
            if (blobCategorical != null)
            {
                m_colCategoricalTop.Clear();
                for (int i = 0; i < m_param.categorical_trans_param.num_input; i++)
                {
                    Blob<T> blobTop = new Blob<T>(m_cuda, m_log);
                    m_colCategoricalTop.Add(blobTop);
                }
 
                LayerParameter p = new LayerParameter(LayerParameter.LayerType.CATEGORICAL_TRANS, m_param.name + ".categorical");
                p.categorical_trans_param.Copy(m_param.categorical_trans_param);
                m_categoricalLayer = Layer<T>.Create(m_cuda, m_log, convertLayerParam(p, m_param), null) as CategoricalTransformationLayer<T>;
 
                m_colBtm.Clear();
                m_colBtm.Add(blobCategorical);
                m_categoricalLayer.LayerSetUp(m_colBtm, m_colCategoricalTop);
                blobs.Add(m_categoricalLayer.blobs);
            }
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            Blob<T> blobNumeric = null;
            Blob<T> blobCategorical = null;
            int nN = 0;
            int nC = 0;
            int? nH = null;
 
            getBlobs(colBottom, out blobNumeric, out blobCategorical);
 
            if (blobNumeric != null)
            {
                nN = colBottom[0].num;
                m_colBtm.Clear();
                m_colBtm.Add(blobNumeric);
                m_numericLayer.Reshape(m_colBtm, m_colNumericTop);
 
                if (colBottom[0].num_axes > 2)
                {
                    nC = colBottom[0].channels;
                    nH = (int)m_param.numeric_trans_param.num_input;
                }
                else
                {
                    nC += (int)m_param.numeric_trans_param.num_input;
                }
            }
 
            if (blobCategorical != null)
            {
                if (nN != 0)
                    m_log.CHECK_EQ(colBottom[1].num, nN, "The bottom(0).num and bottom(1).num must be equal!");
 
                nN = colBottom[1].num;
                m_colBtm.Clear();
                m_colBtm.Add(blobCategorical);
                m_categoricalLayer.Reshape(m_colBtm, m_colCategoricalTop);
 
                if (colBottom[0].num_axes > 2)
                {
                    if (nC != 0)
                        m_log.CHECK_EQ(colBottom[1].channels, nC, "The bottom(0).channels and bottom(1).channels must be equal!");
 
                    nC = colBottom[1].channels;
 
                    if (!nH.HasValue)
                        nH = 0;
                    nH += (int)m_param.categorical_trans_param.num_input;
                }
                else
                {
                    nC += (int)m_param.categorical_trans_param.num_input;
                }
            }
 
            int nEmb = (int)m_param.categorical_trans_param.state_size;
 
            m_rgShape.Clear();
            m_rgShape.Add(nN);
 
            if (!nH.HasValue)
            {
                m_rgShape.Add(nC * nEmb);
            }
            else
            {
                m_rgShape.Add(nC);
                m_rgShape.Add(nH.Value * nEmb);
            }
 
            colTop[0].Reshape(m_rgShape);
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            Blob<T> blobNumeric = null;
            Blob<T> blobCategorical = null;
            int nCount = 0;
            int nBlocks = (int)m_param.numeric_trans_param.num_input + (int)m_param.categorical_trans_param.num_input;
            int nEmb = (int)m_param.numeric_trans_param.state_size;
            int nIdx = 0;
 
            getBlobs(colBottom, out blobNumeric, out blobCategorical);
 
            if (blobNumeric != null)
            {
                m_colBtm.Clear();
                m_colBtm.Add(blobNumeric);
                m_numericLayer.Forward(m_colBtm, m_colNumericTop);
                nCount = m_colNumericTop[0].count();
 
                for (int i = 0; i < m_colNumericTop.Count; i++)
                {
                    m_cuda.channel_copy(nCount, m_colNumericTop[0].num, 1, nBlocks, nEmb, nIdx, colTop[0].mutable_gpu_data, m_colNumericTop[i].gpu_data, DIR.BWD);
                    nIdx++;
                }
            }
 
            if (blobCategorical != null)
            {
                m_colBtm.Clear();
                m_colBtm.Add(blobCategorical);
                m_categoricalLayer.Forward(m_colBtm, m_colCategoricalTop);
                nCount = m_colCategoricalTop[0].count();
 
                for (int i = 0; i < m_colCategoricalTop.Count; i++)
                {
                    m_cuda.channel_copy(nCount, m_colCategoricalTop[0].num, 1, nBlocks, nEmb, nIdx, colTop[0].mutable_gpu_data, m_colCategoricalTop[i].gpu_data, DIR.BWD);
                    nIdx++;
                }
            }
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            Blob<T> blobNumeric = null;
            Blob<T> blobCategorical = null;
            int nCount;
            int nBlocks = (int)m_param.numeric_trans_param.num_input + (int)m_param.categorical_trans_param.num_input;
            int nEmb = (int)m_param.numeric_trans_param.state_size;
            int nIdx = 0;
 
            getBlobs(colBottom, out blobNumeric, out blobCategorical);
 
            if (blobNumeric != null)
            {
                nCount = m_colNumericTop[0].count();
 
                for (int i = 0; i < m_colNumericTop.Count; i++)
                {
                    m_cuda.channel_copy(nCount, m_colNumericTop[0].num, 1, nBlocks, nEmb, nIdx, colTop[0].gpu_diff, m_colNumericTop[i].mutable_gpu_diff, DIR.FWD);
                    nIdx++;
                }
 
                m_colBtm.Clear();
                m_colBtm.Add(blobNumeric);
                m_numericLayer.Backward(m_colNumericTop, new List<bool>() { true }, m_colBtm);
            }
 
            if (blobCategorical != null)
            {
                nCount = m_colCategoricalTop[0].count();
 
                for (int i = 0; i < m_colCategoricalTop.Count; i++)
                {
                    m_cuda.channel_copy(nCount, m_colCategoricalTop[0].num, 1, nBlocks, nEmb, nIdx, colTop[0].gpu_diff, m_colCategoricalTop[i].mutable_gpu_diff, DIR.FWD);
                    nIdx++;
                }
 
                m_colBtm.Clear();
                m_colBtm.Add(blobCategorical);
                m_categoricalLayer.Backward(m_colCategoricalTop, new List<bool>() { true }, m_colBtm);
            }
        }
    }
}