MyCaffeTrainerRNN.cs

using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Diagnostics;
using System.Linq;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using MyCaffe.basecode;
using MyCaffe.basecode.descriptors;
using MyCaffe.common;
using MyCaffe.gym;
using MyCaffe.param;
 
namespace MyCaffe.trainers
{
    public partial class MyCaffeTrainerRNN : Component, IXMyCaffeCustomTrainerRNN, IxTrainerCallbackRNN
    {
        protected PropertySet m_properties = null;
        protected int m_nProjectID = 0;
        protected ConnectInfo m_dsCi = null;
        IxTrainerRNN m_itrainer = null;
        TRAINER_TYPE m_trainerType = TRAINER_TYPE.RNN_SIMPLE;
        IXMyCaffeCustomTrainerCallback m_icallback = null;
        CryptoRandom m_random = new CryptoRandom();
        int m_nSnapshot = 0;
        bool m_bSnapshot = false;
        double m_dfLoss = 0;
        double m_dfAccuracy = 0;
        int m_nIteration = 0;
        int m_nIterations = -1;
        BucketCollection m_rgVocabulary = null;
        object m_syncObj = new object();
 
        enum TRAINER_TYPE
        {
            RNN_SIMPLE,
            RNN_SUPER_SIMPLE
        }
 
        public MyCaffeTrainerRNN()
        {
            InitializeComponent();
        }
 
        public MyCaffeTrainerRNN(IContainer container)
        {
            container.Add(this);
 
            InitializeComponent();
        }
 
        #region Overrides
 
        protected virtual string name
        {
            get { return "MyCaffe RNN Trainer"; }
        }
 
        protected virtual TRAINING_CATEGORY category
        {
            get { return TRAINING_CATEGORY.RECURRENT; }
        }
 
        protected virtual DatasetDescriptor get_dataset_override(int nProjectID, ConnectInfo ci = null)
        {
            return null;
        }
 
        protected virtual string get_information()
        {
            return "";
        }
 
        protected virtual IxTrainerRNN create_trainerD(Component caffe)
        {
            MyCaffeControl<double> mycaffe = caffe as MyCaffeControl<double>;
            m_nProjectID = mycaffe.CurrentProject.OriginalID;
            m_dsCi = mycaffe.DatasetConnectInfo;
 
            int.TryParse(mycaffe.CurrentProject.GetSolverSetting("max_iter"), out m_nIterations);
            int.TryParse(mycaffe.CurrentProject.GetSolverSetting("snapshot"), out m_nSnapshot);
 
            switch (m_trainerType)
            {
                case TRAINER_TYPE.RNN_SUPER_SIMPLE:
                    return new rnn.simple.TrainerRNNSimple<double>(mycaffe, m_properties, m_random, this, m_rgVocabulary);
 
                case TRAINER_TYPE.RNN_SIMPLE:
                    return new rnn.simple.TrainerRNN<double>(mycaffe, m_properties, m_random, this, m_rgVocabulary);
 
                default:
                    throw new Exception("Unknown trainer type '" + m_trainerType.ToString() + "'!");
            }
        }
 
        protected virtual IxTrainerRNN create_trainerF(Component caffe)
        {
            MyCaffeControl<float> mycaffe = caffe as MyCaffeControl<float>;
            m_nProjectID = mycaffe.CurrentProject.OriginalID;
            m_dsCi = mycaffe.DatasetConnectInfo;
 
            int.TryParse(mycaffe.CurrentProject.GetSolverSetting("max_iter"), out m_nIterations);
            int.TryParse(mycaffe.CurrentProject.GetSolverSetting("snapshot"), out m_nSnapshot);
 
            switch (m_trainerType)
            {
                case TRAINER_TYPE.RNN_SUPER_SIMPLE:
                    return new rnn.simple.TrainerRNNSimple<float>(mycaffe, m_properties, m_random, this, m_rgVocabulary);
 
                case TRAINER_TYPE.RNN_SIMPLE:
                    return new rnn.simple.TrainerRNN<float>(mycaffe, m_properties, m_random, this, m_rgVocabulary);
 
                default:
                    throw new Exception("Unknown trainer type '" + m_trainerType.ToString() + "'!");
            }
        }
 
        protected virtual void dispose()
        {
        }
 
        protected virtual void initialize(InitializeArgs e)
        {
        }
 
        protected virtual void shutdown()
        {
        }
 
        protected virtual bool getData(GetDataArgs e)
        {
            return false;
        }
 
        protected virtual bool convertOutput(ConvertOutputArgs e)
        {
            return false;
        }
 
        protected virtual void testAccuracyUpdate(TestAccuracyUpdateArgs e)
        {
        }
 
        protected virtual bool get_update_snapshot(out int nIteration, out double dfAccuracy)
        {
            nIteration = (int)GetProperty("GlobalIteration");
            dfAccuracy = GetProperty("GlobalAccuracy");
 
            if (m_bSnapshot)
            {
                m_bSnapshot = false;
                return true;
            }
 
            return false;
        }
 
        protected virtual void openUi()
        {
        }
 
        protected virtual BucketCollection preloaddata(Log log, CancelEvent evtCancel, int nProjectID, PropertySet propertyOverride = null, ConnectInfo ci = null)
        {
            return null;
        }
 
        #endregion
 
        #region IXMyCaffeCustomTrainer Interface
 
        public Stage Stage
        {
            get { return Stage.RNN; }
        }
 
        public string Name
        {
            get { return name; }
        }
 
        public TRAINING_CATEGORY TrainingCategory
        {
            get { return category; }
        }
 
        public bool GetUpdateSnapshot(out int nIteration, out double dfAccuracy)
        {
            return get_update_snapshot(out nIteration, out dfAccuracy);
        }
 
        public DatasetDescriptor GetDatasetOverride(int nProjectID, ConnectInfo ci = null)
        {
            return get_dataset_override(nProjectID, ci);
        }
 
        public bool IsTrainingSupported
        {
            get { return true; }
        }
 
        public bool IsTestingSupported
        {
            get { return true; }
        }
 
        public bool IsRunningSupported
        {
            get { return true; }
        }
 
        public void CleanUp()
        {
            cleanup(3000, true);
        }
 
        private void cleanup(int nWait, bool bCallShutdown = false)
        {
            lock (m_syncObj)
            {
                if (m_itrainer != null)
                {
                    m_itrainer.Shutdown(nWait);
                    m_itrainer = null;
                }
 
                if (bCallShutdown)
                    shutdown();
            }
        }
 
        public void Initialize(string strProperties, IXMyCaffeCustomTrainerCallback icallback)
        {
            m_icallback = icallback;
            m_properties = new PropertySet(strProperties);
 
            string strTrainerType = m_properties.GetProperty("TrainerType");
 
            switch (strTrainerType)
            {
                case "RNN.SIMPLE":   // bare bones model
                    m_trainerType = TRAINER_TYPE.RNN_SIMPLE;
                    break;
 
                default:
                    throw new Exception("Unknown trainer type '" + strTrainerType + "'!");
            }
        }
 
        private IxTrainerRNN createTrainer(Component mycaffe)
        {
            IxTrainerRNN itrainer = null;
 
            if (mycaffe is MyCaffeControl<double>)
                itrainer = create_trainerD(mycaffe);
            else
                itrainer = create_trainerF(mycaffe);
 
            itrainer.Initialize();
 
            return itrainer;
        }
 
        public float[] Run(Component mycaffe, int nN)
        {
            if (m_itrainer == null)
                m_itrainer = createTrainer(mycaffe);
 
            PropertySet runProp = null;
            IXMyCaffeCustomTrainerCallbackRNN icallback = m_icallback as IXMyCaffeCustomTrainerCallbackRNN;
            if (icallback != null)
                runProp = icallback.GetRunProperties();
 
            float[] rgResults = m_itrainer.Run(nN, runProp);
            cleanup(0);
 
            return rgResults;
        }
 
        public byte[] Run(Component mycaffe, int nN, out string type)
        {
            if (m_itrainer == null)
                m_itrainer = createTrainer(mycaffe);
 
            PropertySet runProp = null;
            IXMyCaffeCustomTrainerCallbackRNN icallback = m_icallback as IXMyCaffeCustomTrainerCallbackRNN;
            if (icallback != null)
                runProp = icallback.GetRunProperties();
 
            byte[] rgResults = m_itrainer.Run(nN, runProp, out type);
            cleanup(0);
 
            return rgResults;
        }
 
        public void Test(Component mycaffe, int nIterationOverride, ITERATOR_TYPE type = ITERATOR_TYPE.ITERATION)
        {
            if (m_itrainer == null)
                m_itrainer = createTrainer(mycaffe);
 
            if (nIterationOverride == -1)
                nIterationOverride = m_nIterations;
 
            m_itrainer.Test(nIterationOverride, type);
            cleanup(0);
        }
 
        public void Train(Component mycaffe, int nIterationOverride, ITERATOR_TYPE type = ITERATOR_TYPE.ITERATION, TRAIN_STEP step = TRAIN_STEP.NONE)
        {
            if (m_itrainer == null)
                m_itrainer = createTrainer(mycaffe);
 
            if (nIterationOverride == -1)
                nIterationOverride = m_nIterations;
 
            m_itrainer.Train(nIterationOverride, type, step);
            cleanup(0);
        }
 
        #endregion
 
        public void OnInitialize(InitializeArgs e)
        {
            initialize(e);
        }
 
        public void OnShutdown()
        {
            shutdown();
        }
 
        public void OnGetData(GetDataArgs e)
        {
            getData(e);
        }
 
        public void OnConvertOutput(ConvertOutputArgs e)
        {
            convertOutput(e);
        }
 
        public void OnTestAccuracyUpdate(TestAccuracyUpdateArgs e)
        {
            testAccuracyUpdate(e);
        }
 
        public void OnUpdateStatus(GetStatusArgs e)
        {
            if (m_icallback != null)
            {
                m_dfLoss = e.Loss;
                m_nIteration = e.Frames;
                m_nIterations = e.MaxFrames;
                m_dfAccuracy = e.TotalReward;
 
                Dictionary<string, double> rgValues = new Dictionary<string, double>();
                rgValues.Add("GlobalIteration", e.Frames);
                rgValues.Add("GlobalLoss", e.Loss);
                rgValues.Add("LearningRate", e.LearningRate);
                rgValues.Add("GlobalAccuracy", e.TotalReward);
                m_icallback.Update(TrainingCategory, rgValues);
            }
        }
 
        public void OnWait(WaitArgs e)
        {
            Thread.Sleep(e.Wait);
        }
 
        public double GetProperty(string strProp)
        {
            switch (strProp)
            {
                case "GlobalLoss":
                    return m_dfLoss;
 
                case "GlobalAccuracy":
                    return m_dfAccuracy;
 
                case "GlobalIteration":
                    return m_nIteration;
 
                case "GlobalMaxIterations":
                    return m_nIterations;
 
                default:
                    throw new Exception("The property '" + strProp + "' is not supported by the MyCaffeTrainerRNN.");
            }
        }
 
        public string Information
        {
            get { return get_information(); }
        }
 
        public void OpenUi()
        {
            openUi();
        }
 
        public BucketCollection PreloadData(Log log, CancelEvent evtCancel, int nProjectID, PropertySet propertyOverride = null, ConnectInfo ci = null)
        {
            return preloaddata(log, evtCancel, nProjectID, propertyOverride, ci);
        }
 
        public string ResizeModel(Log log, string strModel, BucketCollection rgVocabulary)
        {
            if (rgVocabulary == null || rgVocabulary.Count == 0)
                return strModel;
 
            int nVocabCount = rgVocabulary.Count;
            NetParameter p = NetParameter.FromProto(RawProto.Parse(strModel));
            string strEmbedName = "";
            EmbedParameter embed = null;
            string strIpName = "";
            InnerProductParameter ip = null;
 
            foreach (LayerParameter layer in p.layer)
            {
                if (layer.type == LayerParameter.LayerType.EMBED)
                {
                    strEmbedName = layer.name;
                    embed = layer.embed_param;
                }
                else if (layer.type == LayerParameter.LayerType.INNERPRODUCT)
                {
                    strIpName = layer.name;
                    ip = layer.inner_product_param;
                }
            }
 
            if (embed != null)
            {
                if (embed.input_dim != (uint)nVocabCount)
                {
                    log.WriteLine("WARNING: Embed layer '" + strEmbedName + "' input dim changed from " + embed.input_dim.ToString() + " to " + nVocabCount.ToString() + " to accomodate for the vocabulary count.");
                    embed.input_dim = (uint)nVocabCount;
                }
            }
 
            if (ip.num_output != (uint)nVocabCount)
            {
                log.WriteLine("WARNING: InnerProduct layer '" + strIpName + "' num_output changed from " + ip.num_output.ToString() + " to " + nVocabCount.ToString() + " to accomodate for the vocabulary count.");
                ip.num_output = (uint)nVocabCount;
            }
 
            m_rgVocabulary = rgVocabulary;
 
            RawProto proto = p.ToProto("root");
            return proto.ToString();
        }
    }
}