AdamWSolver.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading;
using System.IO;
using MyCaffe.basecode;
using MyCaffe.db.image;
using MyCaffe.common;
using MyCaffe.param;
 
namespace MyCaffe.solvers
{
    public class AdamWSolver<T> : SGDSolver<T>
    {
        double m_dfDetachedWeightDecayRate = 0.0f;
 
        public AdamWSolver(CudaDnn<T> cuda, Log log, SolverParameter p, CancelEvent evtCancel, AutoResetEvent evtForceSnapshot, AutoResetEvent evtForceTest, IXDatabaseBase db, IXPersist<T> persist, int nSolverCount = 1, int nSolverRank = 0, Net<T> shareNet = null, onGetWorkspace getws = null, onSetWorkspace setws = null)
            : base(cuda, log, p, evtCancel, evtForceSnapshot, evtForceTest, db, persist, nSolverCount, nSolverRank, shareNet, getws, setws)
        {
            m_dfDetachedWeightDecayRate = p.adamw_decay;
            AdamPreSolve();
        }
 
        public virtual void AdamPreSolve()
        {
            // Add the extra history entries for AdaDelta after those from
            // SGDSolver::PreSolve
            BlobCollection<T> colNetParams = m_net.learnable_parameters;
 
            for (int i = 0; i < colNetParams.Count; i++)
            {
                List<int> rgShape = colNetParams[i].shape();
                Blob<T> blob = new Blob<T>(m_cuda, m_log, rgShape);
                m_colHistory.Add(blob);
            }
        }
 
        public override void ComputeUpdateValue(int param_id, double dfRate, int nIterationOverride = -1)
        {
            BlobCollection<T> colNetParams = m_net.learnable_parameters;
 
            if (!colNetParams[param_id].DiffExists)
                return;
 
            if (nIterationOverride == -1)
                nIterationOverride = m_nIter;
 
            List<double?> net_params_lr = m_net.params_lr;
            double dfLocalRate = dfRate * net_params_lr[param_id].GetValueOrDefault(0);
            List<double?> net_params_decay = net.params_weight_decay;
            double dfLocalDecay = m_dfDetachedWeightDecayRate * net_params_decay[param_id].GetValueOrDefault(0);
            double dfBeta1 = m_param.momentum;
            T fBeta1 = Utility.ConvertVal<T>(dfBeta1);
            double dfBeta2 = m_param.momentum2;
            T fBeta2 = Utility.ConvertVal<T>(dfBeta2);
 
            // we create aliases for convienience
            int nUpdateHistoryOffset = colNetParams.Count;
            Blob<T> val_m = m_colHistory[param_id];
            Blob<T> val_v = m_colHistory[param_id + nUpdateHistoryOffset];
 
            int nT = nIterationOverride + 1;
            // Set the schedule multiplier
            int nN = colNetParams[param_id].count();
            double dfEpsHat = m_param.delta;
 
            // Compute the update to history, then copy it to the parameter diff.
            m_cuda.adamw_update(nN,
                               colNetParams[param_id].mutable_gpu_diff,
                               val_m.mutable_gpu_data,
                               val_v.mutable_gpu_data,
                               fBeta1,
                               fBeta2,
                               Utility.ConvertVal<T>(dfEpsHat),
                               Utility.ConvertVal<T>(dfLocalRate),
                               Utility.ConvertVal<T>(dfLocalDecay),
                               colNetParams[param_id].mutable_gpu_data,
                               nT);
        }
    }
}