ComputeGraph.cs

using MyCaffe.basecode;
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
 
namespace MyCaffe.common
{
    public class ComputeGraph<T> : IDisposable
    {
        Blob<T> m_blobWork;
        CudaDnn<T> m_cuda;
        Log m_log;
        List<Tuple<string, Action>> m_rgBackprop = new List<Tuple<string, Action>>();
        Dictionary<string, Blob<T>[]> m_rgDebug = new Dictionary<string, Blob<T>[]>();
        bool m_bNeedsBackprop = true;
        bool m_bCheckForNans = false;
        bool m_bClipGradients = false;
        bool m_bAddDebug = false;
        int m_nAxis = 0;
        string m_strMarker = null;
 
        public ComputeGraph(CudaDnn<T> cuda, Log log, int nAxis, bool bNeedsBackprop = true, bool bClipGradients = false, bool bCheckNans = false, bool bAddDebug = false)
        {
            m_cuda = cuda;
            m_log = log;
            m_blobWork = new Blob<T>(cuda, log);
            m_bNeedsBackprop = bNeedsBackprop;
            m_bCheckForNans = bCheckNans;
            m_bClipGradients = bClipGradients;
            m_bAddDebug = bAddDebug;
            m_nAxis = nAxis;            
        }
 
        public void Dispose()
        {
            if (m_blobWork != null)
            {
                m_blobWork.Dispose();
                m_blobWork = null;
            }
        }
 
        private void add_debug(string str, params Blob<T>[] rg)
        {
            string strName = m_rgDebug.Count.ToString() + "_" + str;
            m_rgDebug.Add(strName, rg);
        }
 
        public Dictionary<string, Blob<T>[]> Debug
        {
            get { return m_rgDebug; }
        }
 
        public string marker
        {
            get { return m_strMarker; }
            set { m_strMarker = value; }
        }
 
        public bool needs_backprop
        {
            get { return m_bNeedsBackprop; }
            set { m_bNeedsBackprop = value; }
        }
 
        public int axis
        {
            get { return m_nAxis; }
        }
        
        private Blob<T> work
        {
            get { return m_blobWork; }
        }
 
        private int input_count(Blob<T> b)
        {
            if (b.num_axes <= m_nAxis + 2)
                return 1;
 
            return b.count(m_nAxis + 2);
        }
 
        private void clip_gradient1(Blob<T> b)
        {
            float[] rg = Utility.ConvertVecF<T>(b.mutable_cpu_diff);
            
            for (int i = 0; i < rg.Length; i++)
            {
                if (Math.Abs(rg[i]) < 0.000001)
                    rg[i] = 0;
                else
                    rg[i] = (float)Math.Round(rg[i], 7);
            }
 
            b.mutable_cpu_diff = Utility.ConvertVec<T>(rg);
        }
 
        private void clip_gradient(params Blob<T>[] rg)
        {
            foreach (Blob<T> b in rg)
            {
                clip_gradient1(b);
            }
        }
 
        private T[] round(T[] rgData1, int nDecimals)
        {
            float[] rgData = Utility.ConvertVecF<T>(rgData1);
 
            for (int i = 0; i < rgData.Length; i++)
            {
                rgData[i] = (float)Math.Round(rgData[i], nDecimals);
            }
 
            return Utility.ConvertVec<T>(rgData);
        }
 
        private void check_nan(params Blob<T>[] rg)
        {
            for (int i = 0; i < rg.Length; i++)
            {
                work.ReshapeLike(rg[i]);
                Tuple<double, double, double, double> data = rg[i].minmax_data(work, true);
                Tuple<double, double, double, double> diff = rg[i].minmax_diff(work, true);
 
                double dfDataNanCount = data.Item3;
                double dfDataInfCount = data.Item4;
                double dfDiffNanCount = diff.Item3;
                double dfDiffInfCount = diff.Item4;
 
                if (dfDataNanCount > 0 || dfDataInfCount > 0)
                    throw new Exception("NAN or INF detected in " + rg[i].Name + " data!");
 
                if (dfDataNanCount > 0 || dfDataInfCount > 0)
                    throw new Exception("NAN or INF detected in " + rg[i].Name + " diff!");
            }
        }
 
        private void apply(Blob<T> work, Blob<T> btm)
        {
            m_cuda.add(btm.count(), work.gpu_diff, btm.gpu_diff, btm.mutable_gpu_diff);
        }
 
        public void DebugOp(params Blob<T>[] rgB)
        {
            string strMarker = marker;
            Action backward = () =>
            {
                string str = "";
                for (int i = 0; i < rgB.Length; i++)
                {
                    str += rgB[i].Name + ",";
                }
                str = str.TrimEnd(',');
 
                Trace.WriteLine("Debugging at " + strMarker + " blobs: " + str);
            };
            m_rgBackprop.Add(new Tuple<string, Action>(m_strMarker, backward));
        }
 
        public Blob<T> Round(Blob<T> b, int nDecimals = 6)
        {
            b.mutable_cpu_data = round(b.mutable_cpu_data, nDecimals);
 
            if (m_bNeedsBackprop)
            {
                Action backward = () =>
                {
                    b.mutable_cpu_diff = round(b.mutable_cpu_diff, nDecimals);
                };
                m_rgBackprop.Add(new Tuple<string, Action>(m_strMarker, backward));
            }
 
            return b;
        }
 
        public Blob<T> PeekRow(Blob<T> btm, Blob<T> top, int ix)
        {
            string strMarker = marker;
            List<int> rgShape = new List<int>() { 1, 1 };
            rgShape[1] = btm.count(m_nAxis + 1);
            top.Reshape(rgShape);
 
            int nSpatialDim = btm.count(m_nAxis + 1);
            m_cuda.copy(nSpatialDim, btm.gpu_data, top.mutable_gpu_data, nSpatialDim * ix, 0);
 
            if (m_bNeedsBackprop)
            {
                Action backward = () =>
                {
                    m_cuda.copy(nSpatialDim, top.gpu_diff, btm.mutable_gpu_diff, 0, nSpatialDim * ix);
                    m_cuda.copy(nSpatialDim, top.gpu_data, btm.mutable_gpu_data, 0, nSpatialDim * ix);
 
                    if (m_bCheckForNans)
                        check_nan(btm);
                };
                m_rgBackprop.Add(new Tuple<string, Action>(m_strMarker, backward));
            }
 
            return top;
        }
 
        public void PeekItem(Blob<T> btm, Blob<T> top, int ix)
        {
            string strMarker = marker;
            int nSpatialDim = btm.count(m_nAxis);
 
            m_cuda.copy(nSpatialDim, btm.gpu_data, top.mutable_gpu_data, 0, nSpatialDim * ix);
 
            if (m_bNeedsBackprop)
            {
                Action backward = () =>
                {
                    btm.Reshape(1, 1, 1, 1);
                    m_cuda.copy(nSpatialDim, top.gpu_diff, btm.mutable_gpu_diff, nSpatialDim * ix, 0);
                };
                m_rgBackprop.Add(new Tuple<string, Action>(m_strMarker, backward));
            }
        }
 
        public Blob<T> CopyToRow(Blob<T> btm, Blob<T> top, int ix, bool bCopyDiff = false)
        {
            string strMarker = marker;
            int nSpatialDim = btm.count(m_nAxis);
 
            if (btm.count() == 0)
                top.SetData(0, nSpatialDim * ix, nSpatialDim);
            else
                m_cuda.copy(nSpatialDim, btm.gpu_data, top.mutable_gpu_data, 0, nSpatialDim * ix);
 
            if (m_bNeedsBackprop)
            {
                Action backward = () =>
                {
                    if (bCopyDiff)
                        m_cuda.copy(nSpatialDim, top.gpu_diff, btm.mutable_gpu_diff, nSpatialDim * ix, 0);
                    m_cuda.copy(nSpatialDim, top.gpu_data, btm.mutable_gpu_data, nSpatialDim * ix, 0);
 
                    if (m_bCheckForNans)
                        check_nan(btm);
                };
                m_rgBackprop.Add(new Tuple<string, Action>(m_strMarker, backward));
            }
 
            return top;
        }
 
        public void CopyToCache(Blob<T> btm, Cache<T> cache)
        {
            string strMarker = marker;
            cache.CopyToCache(btm, m_nAxis);
 
            if (m_bNeedsBackprop)
            {
                Action backward = () =>
                {
                    cache.CopyFromCache(btm, m_nAxis);
                };
                m_rgBackprop.Add(new Tuple<string, Action>(m_strMarker, backward));
            }
        }
 
        public Blob<T> tanh(Blob<T> btm, Blob<T> top)
        {
            string strMarker = marker;
            top.ReshapeLike(btm);
 
            m_cuda.tanh_fwd(btm.count(), btm.gpu_data, top.mutable_gpu_data);
 
            if (m_bNeedsBackprop)
            {
                Action backward = () =>
                {
                    work.ReshapeLike(btm);
                    m_cuda.tanh_bwd(top.count(), top.gpu_diff, top.gpu_data, work.mutable_gpu_diff);
                    apply(work, btm);
 
                    if (m_bClipGradients)
                        clip_gradient(btm);
                    if (m_bCheckForNans)
                        check_nan(btm);
                    if (m_bAddDebug)
                        add_debug(strMarker + " - tanh", btm, top);
                };
                m_rgBackprop.Add(new Tuple<string, Action>(m_strMarker, backward));
            }
 
            return top;
        }
 
        public Blob<T> sigmoid(Blob<T> btm, Blob<T> top)
        {
            string strMarker = marker;
            top.ReshapeLike(btm);
 
            m_cuda.sigmoid_fwd(btm.count(), btm.gpu_data, top.mutable_gpu_data);
 
            if (m_bNeedsBackprop)
            {
                Action backward = () =>
                {
                    work.ReshapeLike(btm);
                    m_cuda.sigmoid_bwd(top.count(), top.gpu_diff, top.gpu_data, work.mutable_gpu_diff);
                    apply(work, btm);
 
                    if (m_bClipGradients)
                        clip_gradient(btm);
                    if (m_bCheckForNans)
                        check_nan(btm);
                    if (m_bAddDebug)
                        add_debug(strMarker + " - sigmoid", btm, top);
                };
                m_rgBackprop.Add(new Tuple<string, Action>(m_strMarker, backward));
            }
 
            return top;
        }
 
        public Blob<T> eltmul(Blob<T> btm1, Blob<T> btm2, Blob<T> top)
        {
            string strMarker = marker;
            top.ReshapeLike(btm1);
 
            m_cuda.mul(top.count(), btm1.gpu_data, btm2.gpu_data, top.mutable_gpu_data);
 
            if (m_bNeedsBackprop)
            {
                Action backward = () =>
                {
                    work.ReshapeLike(btm1);
                    m_cuda.mul(btm2.count(), btm2.gpu_data, top.gpu_diff, work.mutable_gpu_diff);
                    apply(work, btm1);
                    work.ReshapeLike(btm2);
                    m_cuda.mul(btm1.count(), btm1.gpu_data, top.gpu_diff, work.mutable_gpu_diff);
                    apply(work, btm2);
 
                    if (m_bClipGradients)
                        clip_gradient(btm1, btm2);
                    if (m_bCheckForNans)
                        check_nan(btm1, btm2);
                    if (m_bAddDebug)
                        add_debug(strMarker + " - eltmul", btm1, btm2, top);
                };
                m_rgBackprop.Add(new Tuple<string, Action>(m_strMarker, backward));
            }
 
            return top;
        }
 
        public Blob<T> scalemul(Blob<T> btm1, Blob<T> btm2, Blob<T> top, int nIdx = 0)
        {
            string strMarker = marker;
            top.ReshapeLike(btm1);
 
            T fScale = btm2.GetData(nIdx);
            m_cuda.scale(top.count(), fScale, btm1.gpu_data, top.mutable_gpu_data);
 
            if (m_bNeedsBackprop)
            {
                Action backward = () =>
                {
                    work.ReshapeLike(btm1);
                    m_cuda.scale(top.count(), fScale, top.gpu_diff, work.mutable_gpu_diff);
                    apply(work, btm1);
 
                    work.ReshapeLike(btm2);
                    float fDot = m_cuda.dot_float(btm1.count(), btm1.gpu_data, top.gpu_diff);
                    work.SetDiff(0);
                    work.SetDiff(fDot, nIdx);
                    apply(work, btm2);
 
                    if (m_bClipGradients)
                        clip_gradient(btm1, btm2);
                    if (m_bCheckForNans)
                        check_nan(btm1, btm2);
                    if (m_bAddDebug)
                        add_debug(strMarker + " - scalemul", btm1, btm2, top);
                };
                m_rgBackprop.Add(new Tuple<string, Action>(m_strMarker, backward));
            }
 
            return top;
        }
 
        public Blob<T> mul(Blob<T> btm1, Blob<T> btm2, Blob<T> top, bool bAccumulateGrad = true)
        {
            string strMarker = marker;
            int nM = btm1.shape(m_nAxis);
            int nN = btm2.count(m_nAxis + 1);
            int nK = btm1.count(m_nAxis + 1);
 
            List<int> rgShape = Utility.Create<int>(m_nAxis, 1);
            rgShape.Add(nM);
            rgShape.Add(nN);
 
            top.Reshape(rgShape);
 
            m_cuda.gemm(false, false, nM, nN, nK, Blob<T>.One, btm1.gpu_data, btm2.gpu_data, Blob<T>.Zero, top.mutable_gpu_data);
 
            if (m_bNeedsBackprop)
            {
                Action backward = () =>
                {
                    T fBeta = (bAccumulateGrad) ? Blob<T>.One : Blob<T>.Zero;
                    m_cuda.gemm(false, true, nM, nK, nN, Blob<T>.One, top.gpu_diff, btm2.gpu_data, fBeta, btm1.mutable_gpu_diff);
                    m_cuda.gemm(true, false, nK, nN, nM, Blob<T>.One, btm1.gpu_data, top.gpu_diff, fBeta, btm2.mutable_gpu_diff);
 
                    if (m_bClipGradients)
                        clip_gradient(btm1, btm2);
                    if (m_bCheckForNans)
                        check_nan(btm1, btm2);
                    if (m_bAddDebug)
                        add_debug(strMarker + " - mul", btm1, btm2, top);
                };
                m_rgBackprop.Add(new Tuple<string, Action>(m_strMarker, backward));
            }
 
            return top;
        }
 
        public Blob<T> add(Blob<T> btm1, Blob<T> btm2, Blob<T> top, bool bAccumulateGrad = true)
        {
            string strMarker = marker;
            top.ReshapeLike(btm1);
 
            m_cuda.add(top.count(), btm1.gpu_data, btm2.gpu_data, top.mutable_gpu_data);
 
            if (m_bNeedsBackprop)
            {
                Action backward = () =>
                {
                    if (!bAccumulateGrad)
                    {
                        btm1.SetDiff(0);
                        btm2.SetDiff(0);
                    }
 
                    m_cuda.add(btm1.count(), btm1.gpu_diff, top.gpu_diff, btm1.mutable_gpu_diff);
                    m_cuda.add(btm2.count(), btm2.gpu_diff, top.gpu_diff, btm2.mutable_gpu_diff);
 
                    if (m_bClipGradients)
                        clip_gradient(btm1, btm2);
                    if (m_bCheckForNans)
                        check_nan(btm1, btm2);
                    if (m_bAddDebug)
                        add_debug(strMarker + " - add", btm1, btm2, top);
                };
                m_rgBackprop.Add(new Tuple<string, Action>(m_strMarker, backward));
            }
 
            return top;
        }
 
        public Blob<T> clear_grad(Blob<T> b)
        {
            if (m_bNeedsBackprop)
            {
                Action backward = () =>
                {
                    b.SetDiff(0);
                };
                m_rgBackprop.Add(new Tuple<string, Action>(m_strMarker, backward));
            }
 
            return b;
        }
 
        public void clear_grad(BlobCollection<T> rg)
        {
            foreach (Blob<T> b in rg)
            {
                clear_grad(b);
            }
        }
 
        public Blob<T> softmax(Blob<T> btm, Blob<T> top)
        {
            string strMarker = marker;
            top.ReshapeLike(btm);
 
            int nOuterNum = btm.count(0, m_nAxis);
            int nInnerNum = btm.count(m_nAxis + 1);
            int nChannels = top.shape(m_nAxis);
            int nCount = btm.count();
 
            work.ReshapeLike(top);
 
            m_cuda.copy(nCount, btm.gpu_data, top.mutable_gpu_data);
 
            // We need to subtract the max to avoid numerical issues, compute the exp
            // and then normalize.
            // compute max.
            m_cuda.channel_max(nOuterNum * nInnerNum, nOuterNum, nChannels, nInnerNum, top.gpu_data, work.mutable_gpu_data);
            // subtract
            m_cuda.channel_sub(nCount, nOuterNum, nChannels, nInnerNum, work.gpu_data, top.mutable_gpu_data);
            // exponentiate
            m_cuda.exp(nCount, top.gpu_data, top.mutable_gpu_data);
            // Sum after exp
            m_cuda.channel_sum(nOuterNum * nInnerNum, nOuterNum, nChannels, nInnerNum, top.gpu_data, work.mutable_gpu_data);
            // divide
            m_cuda.channel_div(nCount, nOuterNum, nChannels, nInnerNum, work.gpu_data, top.mutable_gpu_data);
 
            if (m_bNeedsBackprop)
            {
                Action backward = () =>
                {
                    work.ReshapeLike(top);
                    m_cuda.copy(nCount, top.gpu_diff, work.mutable_gpu_diff);
 
                    // Compute inner1d(top_diff, top_data) and subtract them from the bottom diff.
                    m_cuda.channel_dot(nOuterNum * nInnerNum, nOuterNum, nChannels, nInnerNum, top.gpu_diff, top.gpu_data, work.mutable_gpu_data);
                    m_cuda.channel_sub(nCount, nOuterNum, nChannels, nInnerNum, work.gpu_data, work.mutable_gpu_diff);
 
                    // elementwise multiplication
                    m_cuda.mul(nCount, work.gpu_diff, top.gpu_data, work.mutable_gpu_diff);
                    apply(work, btm);
 
                    if (m_bClipGradients)
                        clip_gradient(btm);
                    if (m_bCheckForNans)
                        check_nan(btm);
                    if (m_bAddDebug)
                        add_debug(strMarker + " - softmax", btm, top);
                };
                m_rgBackprop.Add(new Tuple<string, Action>(m_strMarker, backward));
            }
 
            return top;
        }
 
        public int BackwardCount
        {
            get { return m_rgBackprop.Count; }
        }
 
        public void BackwardOne(int nIdx)
        {
            m_rgBackprop[nIdx].Item2();
        }
 
        public void Backward(bool bClear = false)
        {
            for (int i = m_rgBackprop.Count - 1; i >= 0; i--)
            {
                m_rgBackprop[i].Item2();
            }
 
            if (bClear)
                m_rgBackprop.Clear();
        }
 
        public void Clear()
        {
            m_rgBackprop.Clear();
        }
    }
 
    public class Cache<T> : IDisposable
    {
        CudaDnn<T> m_cuda;
        Blob<T> m_blobCache;
        int m_nCacheIdx = 0;
 
        public Cache(CudaDnn<T> cuda, Log log)
        {
            m_cuda = cuda;
            m_blobCache = new Blob<T>(cuda, log, false);
            m_blobCache.Name = "cache";
        }
 
        public void Dispose()
        {
            if (m_blobCache != null)
            {
                m_blobCache.Dispose();
                m_blobCache = null;
            }
        }
 
        public void Create(int nCount, List<int> rgItemShape)
        {
            List<int> rgShape = new List<int>(rgItemShape);
            rgShape.Insert(0, nCount);
            m_blobCache.Reshape(rgShape);
        }
 
        public void Reset()
        {
            m_nCacheIdx = 0;
            m_blobCache.SetData(0);
        }
 
        public void CopyToCache(Blob<T> b, int nAxis)
        {
            int nSpatialDim = b.count(nAxis);
 
            if (m_nCacheIdx >= m_blobCache.num)
                throw new Exception("The cache is full!");
 
            m_cuda.copy(nSpatialDim, b.gpu_data, m_blobCache.mutable_gpu_data, 0, nSpatialDim * m_nCacheIdx);
            m_nCacheIdx++;
        }
 
        public void CopyFromCache(Blob<T> b, int nAxis)
        {
            int nSpatialDim = b.count(nAxis);
 
            m_nCacheIdx--;
            if (m_nCacheIdx < 0)
                throw new Exception("The cache is empty!");
 
            m_cuda.copy(nSpatialDim, m_blobCache.gpu_data, b.mutable_gpu_data, nSpatialDim * m_nCacheIdx, 0);
        }
    }
}