EmbedLayer.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using MyCaffe.basecode;
using MyCaffe.common;
using MyCaffe.param;
using MyCaffe.fillers;
using System.Diagnostics;
 
namespace MyCaffe.layers
{
    public class EmbedLayer<T> : Layer<T>
    {
        int m_nM;
        int m_nK;
        int m_nN;
        bool m_bBiasTerm;
        Blob<T> m_blobBiasMultiplier;
        bool m_bWarningShown = false;
#if DEBUG
        Blob<T> m_blobWork;
#endif
 
        public EmbedLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.EMBED;
            m_blobBiasMultiplier = new common.Blob<T>(cuda, log);
            m_blobBiasMultiplier.Name = m_param.name + " biasmult";
 
#if DEBUG
            m_blobWork = new common.Blob<T>(cuda, log);
            m_blobWork.Name = m_param.name + " work";
#endif
            setup_internal_blobs(m_colInternalBlobs);
        }
 
        protected override void dispose()
        {
            if (m_blobBiasMultiplier != null)
            {
                m_blobBiasMultiplier.Dispose();
                m_blobBiasMultiplier = null;
            }
 
#if DEBUG
            if (m_blobWork != null)
            {
                m_blobWork.Dispose();
                m_blobWork = null;
            }
#endif
 
            base.dispose();
        }
 
        protected override void setup_internal_blobs(BlobCollection<T> col)
        {
            if (col.Count > 0)
                return;
 
            col.Add(m_blobBiasMultiplier);
        }
 
        public override int MinBottomBlobs
        {
            get { return 1; }
        }
 
        public override int MaxBottomBlobs
        {
            get { return 2; }
        }
 
        public override int ExactNumTopBlobs
        {
            get { return 1; }
        }
 
        public override bool ReInitializeParameters(WEIGHT_TARGET target)
        {
            base.ReInitializeParameters(target);
 
            if (target == WEIGHT_TARGET.BOTH || target == WEIGHT_TARGET.WEIGHTS)
            {
                Filler<T> weight_filler = Filler<T>.Create(m_cuda, m_log, m_param.embed_param.weight_filler);
                weight_filler.Fill(m_colBlobs[0]);
            }
 
            if (m_param.embed_param.bias_term && m_colBlobs.Count > 1 && (target == WEIGHT_TARGET.BOTH || target == WEIGHT_TARGET.BIAS))
            {
                Filler<T> bias_filler = Filler<T>.Create(m_cuda, m_log, m_param.embed_param.bias_filler);
                bias_filler.Fill(m_colBlobs[1]);
            }
 
            return true;
        }
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            m_bWarningShown = false;
 
            if (colBottom.Count > 1)
                m_param.embed_param.input_dim = (uint)convertF(colBottom[1].GetData(0));
 
            m_nN = (int)m_param.embed_param.num_output;
            m_log.CHECK_GT(m_nN, 0, "EmbedLayer num_output must be positive.");
 
            m_nK = (int)m_param.embed_param.input_dim;
            m_log.CHECK_GT(m_nK, 0, "EmbedLayer input_dim must be positive.");
 
            m_bBiasTerm = m_param.embed_param.bias_term;
 
            if (m_colBlobs.Count > 0)
            {
                m_log.WriteLine("Skipping parameter initialization.");
            }
            else
            {
                m_colBlobs.Clear();
 
                // Initialize the weights --
                // transposed from InnerProductLayer for spacial locality.
                List<int> rgWeightShape = new List<int>() { m_nK, m_nN };
                Blob<T> blobWeight = new Blob<T>(m_cuda, m_log);
                blobWeight.Name = m_param.name + " weights";
                blobWeight.type = BLOB_TYPE.WEIGHT;
 
                if (!shareParameter(blobWeight, rgWeightShape))
                {
                    blobWeight.Reshape(rgWeightShape);
 
                    // fill the weights
                    Filler<T> weight_filler = Filler<T>.Create(m_cuda, m_log, m_param.embed_param.weight_filler);
                    weight_filler.Fill(blobWeight);
                }
                m_colBlobs.Add(blobWeight);
 
 
                // If necessary, initialize and fill the bias term
                if (m_bBiasTerm)
                {
                    List<int> rgBiasShape = new List<int>() { m_nN };
                    Blob<T> blobBias = new Blob<T>(m_cuda, m_log);
                    blobBias.Name = m_param.name + " bias";
                    blobBias.type = BLOB_TYPE.WEIGHT;
 
                    if (!shareParameter(blobBias, rgBiasShape))
                    {
                        blobBias.Reshape(rgBiasShape);
                        Filler<T> bias_filler = Filler<T>.Create(m_cuda, m_log, m_param.embed_param.bias_filler);
                        bias_filler.Fill(blobBias);
                    }
                    m_colBlobs.Add(blobBias);
                }
            }
 
            m_rgbParamPropagateDown = new DictionaryMap<bool>(m_colBlobs.Count, true);
        }
 
        public override void Reshape(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            // Figure out the dimensions
            m_nM = colBottom[0].count();
            List<int> rgTopShape = Utility.Clone<int>(colBottom[0].shape());
            
            rgTopShape.Add(m_nN);
            colTop[0].Reshape(rgTopShape);
 
            // Set up the bias multiplier
            if (m_bBiasTerm)
            {
                List<int> rgBiasShape = new List<int>() { m_nM };
                shareLayerBlob(m_blobBiasMultiplier, rgBiasShape);
                m_blobBiasMultiplier.Reshape(rgBiasShape);
                m_blobBiasMultiplier.SetData(1.0);
            }
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            long hBottomData = colBottom[0].gpu_data;
            long hTopData = colTop[0].mutable_gpu_data;
            long hWeight = m_colBlobs[0].gpu_data;
            int nCount = colTop[0].count();
 
#if DEBUG
            Tuple<double, double, double, double> minmax = colBottom[0].minmax_data(m_blobWork, true);
            double dfMin = minmax.Item1;
            double dfMax = minmax.Item2;
            if (dfMin < 0 || dfMax >= m_nK)
                throw new Exception("A data element within '" + colBottom[0].Name + "' is out of range [0," + m_nK.ToString() + ") non inclusive.  Data Min = " + dfMin.ToString() + " Max = " + dfMax.ToString() + ".");
#endif
 
            m_cuda.embed_fwd(nCount, hBottomData, hWeight, m_nM, m_nN, m_nK, hTopData);
 
            if (m_bBiasTerm)
                m_cuda.gemm(false, false, m_nM, m_nN, 1, 1.0, m_blobBiasMultiplier.gpu_data, m_colBlobs[1].gpu_data, 1.0, hTopData);
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            if (rgbPropagateDown[0] && !m_bWarningShown)
            {
                m_log.WriteLine("WARNING: Can't backpropagate to EmbedLayer input.");
                m_bWarningShown = true;
            }
 
            if (m_rgbParamPropagateDown[0])
            {
                int nTopCount = colTop[0].count();
                long hTopDiff = colTop[0].gpu_diff;
                long hBottomData = colBottom[0].gpu_data;
                long hWeightDiff = m_colBlobs[0].mutable_gpu_diff;
                m_cuda.embed_bwd(nTopCount, hBottomData, hTopDiff, m_nM, m_nN, m_nK, hWeightDiff);                
            }
 
            if (m_bBiasTerm && m_rgbParamPropagateDown[1])
            {
                long hTopDiff = colTop[0].gpu_diff;
                long hBiasDiff = m_colBlobs[1].mutable_gpu_diff;
                m_cuda.gemv(true, m_nM, m_nN, 1.0, hTopDiff, m_blobBiasMultiplier.gpu_data, 1.0, hBiasDiff);
            }
        }
    }
}