MyCaffeTrainerDual.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 MyCaffeTrainerDual : Component, IXMyCaffeCustomTrainerRL, IXMyCaffeCustomTrainerRNN, IxTrainerCallbackRNN
    {
        protected CryptoRandom m_random = new CryptoRandom();
        protected PropertySet m_properties = null;
        protected int m_nProjectID = 0;
        protected ConnectInfo m_dsCi = null;
        IxTrainer m_itrainer = null;
        double m_dfExplorationRate = 0;
        double m_dfOptimalSelectionRate = 0;
        double m_dfImmediateRewards = 0;
        double m_dfGlobalRewards = 0;
        double m_dfGlobalRewardsAve = 0;
        double m_dfGlobalRewardsMax = -double.MaxValue;
        int m_nGlobalEpisodeCount = 0;
        int m_nGlobalEpisodeMax = 0;
        double m_dfLoss = 0;
        int m_nThreads = 1;
        REWARD_TYPE m_rewardType = REWARD_TYPE.MAXIMUM;
        TRAINER_TYPE m_trainerType = TRAINER_TYPE.PG_ST;
        double m_dfAccuracy = 0;
        int m_nIteration = 0;
        int m_nIterations = -1;
        IXMyCaffeCustomTrainerCallback m_icallback = null;
        int m_nSnapshot = 0;
        bool m_bSnapshot = false;
        BucketCollection m_rgVocabulary = null;
        Stage m_stage = Stage.RL;
        bool m_bUsePreloadData = false;
        object m_syncObj = new object();
 
        enum TRAINER_TYPE
        {
            PG_MT,
            PG_ST,
            PG_SIMPLE,
            C51_ST,
            C51b_ST,
            DQN_ST,
            DQN_SIMPLE,
            RNN_SIMPLE,
            RNN_SUPER_SIMPLE
        }
 
        enum REWARD_TYPE
        {
            VALUE,
            AVERAGE,
            MAXIMUM
        }
 
        public MyCaffeTrainerDual()
        {
            InitializeComponent();
        }
 
        public MyCaffeTrainerDual(IContainer container)
        {
            container.Add(this);
 
            InitializeComponent();
        }
 
        #region Overrides
 
        protected virtual string name
        {
            get { return "MyCaffe RL/RNN Dual Trainer"; }
        }
 
        protected virtual TRAINING_CATEGORY category
        {
            get { return TRAINING_CATEGORY.DUAL; }
        }
 
        protected virtual DatasetDescriptor get_dataset_override(int nProjectID, ConnectInfo ci = null)
        {
            return null;
        }
 
        protected virtual string get_information()
        {
            return "";
        }
 
        protected virtual IxTrainer create_trainerD(Component caffe, Stage stage)
        {
            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);
 
            m_properties.SetProperty("UsePreLoadData", m_bUsePreloadData.ToString());
 
            if (stage == Stage.RNN)
            {
                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("The trainer type '" + m_trainerType.ToString() + "' is not supported in the RNN stage!");
                }
            }
            else
            {
                switch (m_trainerType)
                {
                    case TRAINER_TYPE.PG_SIMPLE:
                        return new pg.simple.TrainerPG<double>(mycaffe, m_properties, m_random, this);
 
                    case TRAINER_TYPE.PG_ST:
                        return new pg.st.TrainerPG<double>(mycaffe, m_properties, m_random, this);
 
                    case TRAINER_TYPE.PG_MT:
                        return new pg.mt.TrainerPG<double>(mycaffe, m_properties, m_random, this);
 
                    case TRAINER_TYPE.C51_ST:
                        return new dqn.c51.st.TrainerC51<double>(mycaffe, m_properties, m_random, this);
 
                    case TRAINER_TYPE.DQN_ST:
                        return new dqn.noisy.st.TrainerNoisyDqn<double>(mycaffe, m_properties, m_random, this);
 
                    case TRAINER_TYPE.DQN_SIMPLE:
                        return new dqn.noisy.simple.TrainerNoisyDqn<double>(mycaffe, m_properties, m_random, this);
 
                    default:
                        throw new Exception("The trainer type '" + m_trainerType.ToString() + "' is not supported in the RL stage!");
                }
            }
        }
 
        protected virtual IxTrainer create_trainerF(Component caffe, Stage stage)
        {
            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);
 
            m_properties.SetProperty("UsePreLoadData", m_bUsePreloadData.ToString());
 
            if (stage == Stage.RNN)
            {
                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("The trainer type '" + m_trainerType.ToString() + "' is not supported in the RNN stage!");
                }
            }
            else
            {
                switch (m_trainerType)
                {
                    case TRAINER_TYPE.PG_SIMPLE:
                        return new pg.simple.TrainerPG<float>(mycaffe, m_properties, m_random, this);
 
                    case TRAINER_TYPE.PG_ST:
                        return new pg.st.TrainerPG<float>(mycaffe, m_properties, m_random, this);
 
                    case TRAINER_TYPE.PG_MT:
                        return new pg.mt.TrainerPG<float>(mycaffe, m_properties, m_random, this);
 
                    case TRAINER_TYPE.C51_ST:
                        return new dqn.c51.st.TrainerC51<float>(mycaffe, m_properties, m_random, this);
 
                    case TRAINER_TYPE.DQN_ST:
                        return new dqn.noisy.st.TrainerNoisyDqn<float>(mycaffe, m_properties, m_random, this);
 
                    case TRAINER_TYPE.DQN_SIMPLE:
                        return new dqn.noisy.simple.TrainerNoisyDqn<float>(mycaffe, m_properties, m_random, this);
 
                    default:
                        throw new Exception("The trainer type '" + m_trainerType.ToString() + "' is not supported in the RL stage!");
                }
            }
        }
 
        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 = GlobalEpisodeCount;
            dfAccuracy = GlobalRewards;
 
            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 m_stage; }
        }
 
        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)
        {
            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);
            m_nThreads = m_properties.GetPropertyAsInt("Threads", 1);
 
            string strRewardType = m_properties.GetProperty("RewardType", false);
            if (strRewardType == null)
                strRewardType = "VAL";
            else
                strRewardType = strRewardType.ToUpper();
 
            if (strRewardType == "VAL" || strRewardType == "VALUE")
                m_rewardType = REWARD_TYPE.VALUE;
            else if (strRewardType == "AVE" || strRewardType == "AVERAGE")
                m_rewardType = REWARD_TYPE.AVERAGE;
 
            string strTrainerType = m_properties.GetProperty("TrainerType");
 
            switch (strTrainerType)
            {
                case "PG.SIMPLE":   // bare bones model (Sigmoid only)
                    m_trainerType = TRAINER_TYPE.PG_SIMPLE;
                    m_stage = Stage.RL;
                    break;
 
                case "PG.ST":       // single thread (Sigmoid and Softmax)
                    m_trainerType = TRAINER_TYPE.PG_ST;
                    m_stage = Stage.RL;
                    break;
 
                case "PG":
                case "PG.MT":       // multi-thread (Sigmoid and Softmax)
                    m_trainerType = TRAINER_TYPE.PG_MT;
                    m_stage = Stage.RL;
                    break;
 
                case "C51.ST":      // single threaded C51
                    m_trainerType = TRAINER_TYPE.C51_ST;
                    m_stage = Stage.RL;
                    break;
 
                case "DQN.ST": // single threaded Noisy DQN
                    m_trainerType = TRAINER_TYPE.DQN_ST;
                    m_stage = Stage.RL;
                    break;
 
                case "DQN.SIMPLE": // single threaded Noisy DQN
                    m_trainerType = TRAINER_TYPE.DQN_SIMPLE;
                    m_stage = Stage.RL;
                    break;
 
                case "RNN.SIMPLE":
                    m_trainerType = TRAINER_TYPE.RNN_SIMPLE;
                    m_stage = Stage.RNN;
                    break;
 
                case "RNN.SUPER.SIMPLE":
                    m_trainerType = TRAINER_TYPE.RNN_SUPER_SIMPLE;
                    m_stage = Stage.RNN;
                    break;
 
                default:
                    throw new Exception("Unknown trainer type '" + strTrainerType + "'!");
            }
        }
 
        private Stage getStage()
        {
            if (m_trainerType == TRAINER_TYPE.RNN_SIMPLE || m_trainerType == TRAINER_TYPE.RNN_SUPER_SIMPLE)
                return Stage.RNN;
            else
                return Stage.RL;
        }
 
        private IxTrainer createTrainer(Component mycaffe, Stage stage)
        {
            IxTrainer itrainer = null;
 
            if (mycaffe is MyCaffeControl<double>)
                itrainer = create_trainerD(mycaffe, stage);
            else
                itrainer = create_trainerF(mycaffe, stage);
 
            itrainer.Initialize();
 
            return itrainer;
        }
 
        public void Test(Component mycaffe, int nIterationOverride, ITERATOR_TYPE type = ITERATOR_TYPE.ITERATION)
        {
            if (m_itrainer == null)
                m_itrainer = createTrainer(mycaffe, getStage());
 
            if (nIterationOverride == -1)
                nIterationOverride = m_nIterations;
 
            m_itrainer.Test(nIterationOverride, type);
            cleanup(500, false);
        }
 
        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, getStage());
 
            if (nIterationOverride == -1)
                nIterationOverride = m_nIterations;
 
            m_itrainer.Train(nIterationOverride, type, step);
            cleanup(1000, false);
        }
 
        #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)
        {
            m_nIteration = e.Iteration;
            m_dfAccuracy = e.TotalReward;
            m_nIterations = e.MaxFrames;
            m_dfImmediateRewards = e.Reward;
            m_dfGlobalRewards = e.TotalReward;
            m_dfGlobalRewardsMax = Math.Max(m_dfGlobalRewardsMax, e.TotalReward);
            m_dfGlobalRewardsAve = (1.0 / (double)m_nThreads) * e.TotalReward + ((m_nThreads - 1) / (double)m_nThreads) * m_dfGlobalRewardsAve;
            m_dfExplorationRate = e.ExplorationRate;
            m_dfOptimalSelectionRate = e.OptimalSelectionCoefficient;
 
            if (m_nThreads > 1)
                m_nGlobalEpisodeCount++;
            else
                m_nGlobalEpisodeCount = e.Frames;
 
            m_nGlobalEpisodeMax = e.MaxFrames;
            m_dfLoss = e.Loss;
 
            if (m_icallback != null)
            {
                Dictionary<string, double> rgValues = new Dictionary<string, double>();
                rgValues.Add("GlobalIteration", GlobalEpisodeCount);
                rgValues.Add("GlobalLoss", GlobalLoss);
                rgValues.Add("LearningRate", e.LearningRate);
                rgValues.Add("GlobalAccuracy", GlobalRewards);
                rgValues.Add("Threads", m_nThreads);
 
                m_icallback.Update(TrainingCategory, rgValues);
            }
 
            e.NewFrameCount = m_nGlobalEpisodeCount;
 
            if (e.Index == 0 && m_nSnapshot > 0 && m_nGlobalEpisodeCount > 0 && (m_nGlobalEpisodeCount % m_nSnapshot) == 0)
                m_bSnapshot = true;
        }
 
        public void OnWait(WaitArgs e)
        {
            Thread.Sleep(e.Wait);
        }
 
        public double GetProperty(string strProp)
        {
            switch (strProp)
            {
                case "GlobalLoss":
                    return GlobalLoss;
 
                case "GlobalAccuracy": // RNN
                    return m_dfAccuracy;
 
                case "GlobalIteration": // RNN
                    return m_nIteration;
 
                case "GlobalMaxIterations":
                    return m_nIterations;
 
                case "GlobalRewards":
                    return GlobalRewards;
 
                case "GlobalEpisodeCount":
                    return GlobalEpisodeCount;
 
                case "ExplorationRate":
                    return ExplorationRate;
 
                default:
                    throw new Exception("The property '" + strProp + "' is not supported by the MyCaffeTrainerRNN.");
            }
        }
 
        public double GlobalRewards
        {
            get
            {
                switch (m_rewardType)
                {
                    case REWARD_TYPE.VALUE:
                        return m_dfGlobalRewards;
 
                    case REWARD_TYPE.AVERAGE:
                        return m_dfGlobalRewardsAve;
 
                    default:
                        return (m_dfGlobalRewardsMax == -double.MaxValue) ? 0 : m_dfGlobalRewardsMax;
                }
            }
        }
 
        public double ImmediateRewards
        {
            get { return m_dfImmediateRewards; }
        }
 
        public double GlobalLoss
        {
            get { return m_dfLoss; }
        }
 
        public int GlobalIteration
        {
            get { return m_nIteration; }
        }
 
        public int GlobalEpisodeCount
        {
            get { return m_nGlobalEpisodeCount; }
        }
 
        public int GlobalEpisodeMax
        {
            get { return m_nGlobalEpisodeMax; }
        }
 
        public double ExplorationRate
        {
            get { return m_dfExplorationRate; }
        }
 
        public double OptimalSelectionRate
        {
            get { return m_dfOptimalSelectionRate; }
        }
 
        public string Information
        {
            get { return get_information(); }
        }
 
        public void OpenUi()
        {
            openUi();
        }
 
        #region IXMyCaffeCustomTrainerRL Methods
 
        ResultCollection IXMyCaffeCustomTrainerRL.RunOne(Component mycaffe, int nDelay)
        {
            if (m_itrainer == null)
                m_itrainer = createTrainer(mycaffe, Stage.RL);
 
            IxTrainerRL itrainer = m_itrainer as IxTrainerRL;
            if (itrainer == null)
                throw new Exception("The trainer must be set to to 'C51.ST', PG.SIMPLE', 'PG.ST' or 'PG.MT' to run in reinforcement learning mode.");
 
            ResultCollection res = itrainer.RunOne(nDelay);
            cleanup(50, false);
 
            return res;
        }
 
        public byte[] Run(Component mycaffe, int nN, out string type)
        {
            if (m_itrainer == null)
                m_itrainer = createTrainer(mycaffe, Stage.RL);
 
            PropertySet runProp = null;
            IXMyCaffeCustomTrainerCallbackRNN icallback = m_icallback as IXMyCaffeCustomTrainerCallbackRNN;
            if (icallback != null)
                runProp = icallback.GetRunProperties();
 
            IxTrainerRL itrainer = m_itrainer as IxTrainerRL;
            if (itrainer == null)
                throw new Exception("The IxTrainerRL interface must be implemented.");
 
            byte[] rgResults = itrainer.Run(nN, runProp, out type);
            cleanup(0, false);
 
            return rgResults;
        }
 
        #endregion
 
        #region IXMyCaffeCustomTrainerRNN
 
        float[] IXMyCaffeCustomTrainerRNN.Run(Component mycaffe, int nN)
        {
            if (m_itrainer == null)
                m_itrainer = createTrainer(mycaffe, Stage.RNN);
 
            IxTrainerRNN itrainer = m_itrainer as IxTrainerRNN;
            if (itrainer == null)
                throw new Exception("The trainer must be set to to 'RNN.SIMPLE' to run in recurrent learning mode.");
 
            PropertySet runProp = null;
            IXMyCaffeCustomTrainerCallbackRNN icallback = m_icallback as IXMyCaffeCustomTrainerCallbackRNN;
            if (icallback != null)
                runProp = icallback.GetRunProperties();
 
            float[] rgResults = itrainer.Run(nN, runProp);
            cleanup(0, false);
 
            return rgResults;
        }
 
        byte[] IXMyCaffeCustomTrainerRNN.Run(Component mycaffe, int nN, out string type)
        {
            if (m_itrainer == null)
                m_itrainer = createTrainer(mycaffe, Stage.RNN);
 
            IxTrainerRNN itrainer = m_itrainer as IxTrainerRNN;
            if (itrainer == null)
                throw new Exception("The trainer must be set to to 'RNN.SIMPLE' to run in recurrent learning mode.");
 
            PropertySet runProp = null;
            IXMyCaffeCustomTrainerCallbackRNN icallback = m_icallback as IXMyCaffeCustomTrainerCallbackRNN;
            if (icallback != null)
                runProp = icallback.GetRunProperties();
 
            byte[] rgResults = itrainer.Run(nN, runProp, out type);
            m_itrainer.Shutdown(0);
            m_itrainer = null;
 
            return rgResults;
        }
 
        BucketCollection IXMyCaffeCustomTrainerRNN.PreloadData(Log log, CancelEvent evtCancel, int nProjectID, PropertySet propertyOverride = null, ConnectInfo ci = null)
        {
            return preloaddata(log, evtCancel, nProjectID, propertyOverride, ci);
        }
 
        string IXMyCaffeCustomTrainerRNN.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 != null && 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();
        }
 
        #endregion
    }
}