PositionalEncodingLayer.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.gpt
{
    public class PositionalEncodingLayer<T> : Layer<T>
    {
        Blob<T> m_blobPosEnc;
        List<int> m_rgShape = new List<int>() { 1, 1, 1 };
        double m_dfScale;
        int m_nBlockSize;
        int m_nEmbed;
 
        public PositionalEncodingLayer(CudaDnn<T> cuda, Log log, LayerParameter p)
            : base(cuda, log, p)
        {
            m_type = LayerParameter.LayerType.POSITIONAL_ENCODER;
            m_nBlockSize = (int)p.positional_encoder_param.block_size;
            m_nEmbed = (int)p.positional_encoder_param.embed;
            m_dfScale = Math.Sqrt(m_nEmbed);
 
            m_blobPosEnc = new Blob<T>(m_cuda, m_log, false);
            m_blobPosEnc.Name = p.name + " posenc";
 
            setup_internal_blobs(m_colInternalBlobs);
        }
 
        protected override void dispose()
        {
            dispose(ref m_blobPosEnc);
            base.dispose();
        }
 
        protected override void setup_internal_blobs(BlobCollection<T> col)
        {
            if (col.Count > 0)
                return;
 
            col.Add(m_blobPosEnc);
        }
 
        public override int ExactNumBottomBlobs
        {
            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)
        {
            colTop[0].ReshapeLike(colBottom[0]);
 
            int nBatch = colBottom[0].num;
            m_rgShape[0] = nBatch;
            m_rgShape[1] = m_nBlockSize;
            m_rgShape[2] = m_nEmbed;
 
            shareLayerBlob(m_blobPosEnc, m_rgShape);
            if (!m_blobPosEnc.CompareShape(m_rgShape, true))
            {
                m_blobPosEnc.Reshape(m_rgShape);
                m_blobPosEnc.Reshape(1, m_rgShape[1], m_rgShape[2], 1);
                int nDim = m_nBlockSize * m_nEmbed;
 
                if (typeof(T) == typeof(float))
                {
                    float[] rgPosEnc1 = new float[nDim];
                    for (int pos = 0; pos < m_nBlockSize; pos++)
                    {
                        for (int i = 0; i < m_nEmbed; i++)
                        {
                            int nIdx = pos * m_nEmbed + i;
                            double df1 = 2 * i / (double)m_nEmbed;
                            double dfPow = Math.Pow(10000, df1);
                            double dfPos = pos / dfPow;
 
                            if (i % 2 == 0)
                            {
                                double dfSin = Math.Sin(dfPos);
                                rgPosEnc1[nIdx] = (float)dfSin;
                            }
                            else if (i % 2 == 1)
                            {
                                double dfCos = Math.Cos(dfPos);
                                rgPosEnc1[nIdx] = (float)dfCos;
                            }
                        }
                    }
 
                    m_blobPosEnc.mutable_cpu_data = convert(rgPosEnc1);
                }
                else
                {
                    double[] rgPosEnc1 = new double[nDim];
                    for (int pos = 0; pos < m_nBlockSize; pos++)
                    {
                        for (int i = 0; i < m_nEmbed; i++)
                        {
                            int nIdx = pos * m_nEmbed + i;
                            double df1 = 2 * i / (double)m_nEmbed;
                            double dfPow = Math.Pow(10000, df1);
                            double dfPos = pos / dfPow;
 
                            if (i % 2 == 0)
                            {
                                double dfSin = Math.Sin(dfPos);
                                rgPosEnc1[nIdx] = dfSin;
                            }
                            else if (i % 2 == 1)
                            {
                                double dfCos = Math.Cos(dfPos);
                                rgPosEnc1[nIdx] = dfCos;
                            }
                        }
                    }
 
                    m_blobPosEnc.mutable_cpu_data = convert(rgPosEnc1);
                }
 
                if (nBatch > 1)
                {
                    m_blobPosEnc.Reshape(m_rgShape);
 
                    for (int i = 1; i < nBatch; i++)
                    {
                        m_cuda.copy(nDim, m_blobPosEnc.gpu_data, m_blobPosEnc.mutable_gpu_data, 0, i * nDim);
                    }
                }
            }
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            long hBottomData = colBottom[0].gpu_data;
            long hTopData = colTop[0].mutable_gpu_data;
            int nCount = colBottom[0].count();
            
            m_cuda.add(nCount, m_blobPosEnc.gpu_data, hBottomData, hTopData, m_dfScale);
        }
 
        protected override void backward(BlobCollection<T> colTop, List<bool> rgbPropagateDown, BlobCollection<T> colBottom)
        {
            long hTopDiff = colTop[0].gpu_diff;
            long hBottomDiff = colBottom[0].mutable_gpu_diff;
            int nCount = colBottom[0].count();
 
            m_cuda.scale(nCount, m_dfScale, hTopDiff, hBottomDiff);
        }
    }
}