BasePrefetchingDataLayer.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using MyCaffe.basecode;
using MyCaffe.db.image;
using MyCaffe.common;
using MyCaffe.param;
using System.Threading;
using System.Diagnostics;
 
namespace MyCaffe.layers
{
    public abstract class BasePrefetchingDataLayer<T> : BaseDataLayer<T>
    {
        protected Batch<T>[] m_rgPrefetch;
        protected CancelEvent m_evtCancel;
        BlockingQueue<Batch<T>> m_rgPrefetchFree;
        BlockingQueue<Batch<T>> m_rgPrefetchFull;
        Batch<T> m_prefetch_current = null;
        InternalThread<T> m_internalThread;
        Exception m_err = null;
 
        public BasePrefetchingDataLayer(CudaDnn<T> cuda, Log log, LayerParameter p, IXDatabaseBase db, CancelEvent evtCancel)
            : base(cuda, log, p, db)
        {
            m_evtCancel = evtCancel;
 
            m_internalThread = new InternalThread<T>();
            m_internalThread.DoWork += new EventHandler<ActionStateArgs<T>>(m_internalThread_DoWork);
            m_internalThread.OnPreStop += internalThread_OnPreStop;
            m_internalThread.OnPreStart += internalThread_OnPreStart;
 
            m_rgPrefetch = new Batch<T>[p.data_param.prefetch];
            m_rgPrefetchFree = new BlockingQueue<Batch<T>>(m_evtCancel);
            m_rgPrefetchFull = new BlockingQueue<Batch<T>>(m_evtCancel);
 
            for (int i = 0; i < m_rgPrefetch.Length; i++)
            {
                m_rgPrefetch[i] = new Batch<T>(cuda, log);
                m_rgPrefetchFree.Push(m_rgPrefetch[i]);
            }
        }
 
        private void internalThread_OnPreStart(object sender, EventArgs e)
        {
            m_rgPrefetchFree.Reset();
            m_rgPrefetchFull.Reset();
        }
 
        private void internalThread_OnPreStop(object sender, EventArgs e)
        {
            if (m_rgPrefetchFree != null)
                m_rgPrefetchFree.Abort();
 
            if (m_rgPrefetchFull != null)
                m_rgPrefetchFull.Abort();
 
            preStop();
        }
 
        protected virtual void preStop()
        {
        }
 
        protected override void dispose()
        {
            m_internalThread.StopInternalThread();
 
            if (m_rgPrefetchFull != null)
            {
                m_rgPrefetchFull.Dispose();
                m_rgPrefetchFull = null;
            }
 
            if (m_rgPrefetchFree != null)
            {
                m_rgPrefetchFree.Dispose();
                m_rgPrefetchFree = null;
            }
 
            base.dispose();
        }
 
        public override void LayerSetUp(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            base.LayerSetUp(colBottom, colTop);
 
            for (int i = 0; i < m_rgPrefetch.Length; i++)
            {
                m_rgPrefetch[i].Data.update_cpu_data();
 
                if (m_bOutputLabels)
                    m_rgPrefetch[i].Label.update_cpu_data();
            }
 
            m_transformer.InitRand();
 
            if (!delayPrefetch)
            {
                m_log.WriteLine("Initializing prefetch for '" + m_param.name + "'...");
                statupPrefetch();
            }
            else
            {
                m_log.WriteLine("Delaying prefetch for '" + m_param.name + "'...");
            }
        }
 
        protected virtual bool delayPrefetch
        {
            get { return false; }
        }
 
        protected void statupPrefetch()
        {
            m_err = null;
            m_internalThread.StartInternalThread(m_cuda, m_log, m_cuda.GetDeviceID());
            m_log.WriteLine("Prefetch initialized for '" + m_param.name + "'.");
        }
 
        void m_internalThread_DoWork(object sender, ActionStateArgs<T> e)
        {
            CudaDnn<T> cuda = m_cuda;
            Log log = m_log;
            long hStream = 0; // cuda.CreateStream(false);
 
            try
            {
                Batch<T> batch = new Batch<T>(cuda, log);
 
                while (!m_internalThread.CancellationPending)
                {
                    if (m_rgPrefetchFree.Pop(ref batch))
                    {
                        load_batch(batch);
 
                        batch.Data.AsyncGpuPush(hStream);
                        if (hStream != 0)
                            m_cuda.SynchronizeStream(hStream);
 
                        if (m_bOutputLabels)
                        {
                            batch.Label.AsyncGpuPush(hStream);
                            if (hStream != 0)
                                m_cuda.SynchronizeStream(hStream);
                        }
 
                        m_transformer.DistortImage(batch.Data);
 
                        m_rgPrefetchFull.Push(batch);
                    }
                }
            }
            catch (Exception excpt)
            {
                m_err = excpt;
                m_rgPrefetchFull.Abort();
                m_rgPrefetchFree.Abort();
                throw excpt;
            }
            finally
            {
                if (hStream != 0)
                    cuda.FreeStream(hStream);
            }
        }
 
        protected virtual void final_process(Blob<T> blobTop)
        {
            return;
        }
 
        protected override void forward(BlobCollection<T> colBottom, BlobCollection<T> colTop)
        {
            if (m_prefetch_current != null)
                m_rgPrefetchFree.Push(m_prefetch_current);
 
            if (m_rgPrefetchFull.Pop(ref m_prefetch_current))
            {
                // Reshape to loaded data.
                colTop[0].ReshapeLike(m_prefetch_current.Data);
 
                // Copy the data.
                m_cuda.copy(m_prefetch_current.Data.count(), m_prefetch_current.Data.gpu_data, colTop[0].mutable_gpu_data);
                final_process(colTop[0]);
 
                //-----------------------------------------
                // If the blob has a fixed range set, set 
                //  the range in the top data.
                //-----------------------------------------
                m_transformer.SetRange(colTop[0]);
 
                if (m_bOutputLabels)
                {
                    // Reshape to loaded labels.
                    colTop[1].ReshapeLike(m_prefetch_current.Label);
 
                    // Copy the labels.
                    m_cuda.copy(m_prefetch_current.Label.count(), m_prefetch_current.Label.gpu_data, colTop[1].mutable_gpu_data);
                }
            }
            else if (m_err != null)
            {
                throw m_err;
            }
        }
 
        protected abstract void load_batch(Batch<T> batch);
    }
 
    public class Batch<T> : IDisposable 
    {
        Blob<T> m_data;
        Blob<T> m_label;
 
        public Batch(CudaDnn<T> cuda, Log log)
        {
            m_data = new Blob<T>(cuda, log);
            m_label = new Blob<T>(cuda, log);
        }
 
        public void Dispose()
        {
        }
 
        public Blob<T> Data
        {
            get { return m_data; }
        }
 
        public Blob<T> Label
        {
            get { return m_label; }
        }
    }
}