fei-nn-1.0.0: src/MXNet/NN/Optimizer.hs
{-# LANGUAGE CPP #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE UndecidableInstances #-}
module MXNet.NN.Optimizer where
import Control.Lens (use, (.=))
import GHC.Exts (Constraint)
import GHC.TypeLits
import RIO
import qualified RIO.HashMap as M
import RIO.State
import MXNet.Base hiding (Symbol)
import qualified MXNet.Base.Operators.Tensor as T
import MXNet.NN.LrScheduler (LrScheduler (..))
import MXNet.NN.TaggedState (untag)
import MXNet.NN.Types (TaggedModuleState, mod_statistics,
stat_last_lr, stat_num_upd)
-- | Abstract Optimizer type class
class Optimizer (opt :: * -> *) where
data OptimizerTag opt :: *
-- | Specific required arguments
-- data ReqArgs opt :: *
-- | Specific optional arguments
-- type OptArgsList opt :: [KV *]
-- | make the optimizer
makeOptimizer :: (DType dtype, LrScheduler sch, OptimizerCst opt dtype args, MonadIO m)
=> OptimizerTag opt -> sch
-> ArgsHMap (OptimizerSym opt) (NDArray dtype) args
-> m (opt dtype)
-- | run the optimizer with the input & expected tensor
optimize :: (DType dtype, MonadState (TaggedModuleState dtype t) m, MonadIO m)
=> opt dtype -- optimizer
-> Text -- symbol name to optimize
-> NDArray dtype -- parameter
-> NDArray dtype -- gradient
-> m ()
type family OptimizerSym (opt :: * -> *) :: Symbol
type family OptimizerCst (opt :: * -> *) dt (args :: [*]) :: Constraint
-- | SGD optimizer
data SGD_Opt dtype where
SGD_Opt :: (LrScheduler sch, OptimizerCst SGD_Opt dtype args)
=> sch -> ArgsHMap (OptimizerSym SGD_Opt) (NDArray dtype) args
-> SGD_Opt dtype
type instance OptimizerSym SGD_Opt = "_sgd_update"
-- 1.0.0 type instance OptimizerCst SGD_Opt dt args = HasArgs (OptimizerSym SGD_Opt) args '["wd", "rescale_grad", "clip_gradient"]
type instance OptimizerCst SGD_Opt dt args =
HasArgs (OptimizerSym SGD_Opt) (NDArray dt) args '["wd", "rescale_grad", "clip_gradient", "lazy_update"]
instance Optimizer SGD_Opt where
data OptimizerTag SGD_Opt = SGD
makeOptimizer SGD sch args = return $ SGD_Opt sch args
optimize (SGD_Opt sch args) _ weight gradient = do
nup <- use $ untag . mod_statistics . stat_num_upd
let lr = getLR sch nup
untag . mod_statistics . stat_last_lr .= lr
liftIO $ void $ T._sgd_update (#weight := weight
.& #grad := gradient
.& #lr := lr .& args)
(Just [weight])
-- | SGD with momentum optimizer
data SGD_Mom_Opt dtype where
SGD_Mom_Opt :: (LrScheduler sch, OptimizerCst SGD_Mom_Opt dtype args)
=> sch -> ArgsHMap (OptimizerSym SGD_Mom_Opt) (NDArray dtype) args
-> (IORef (M.HashMap Text (NDArray dtype)))
-> SGD_Mom_Opt dtype
type instance OptimizerSym SGD_Mom_Opt = "_sgd_mom_update"
-- 1.0.0 type instance OptimizerCst SGD_Mom_Opt dt args = HasArgs (OptimizerSym SGD_Mom_Opt) args '["momentum", "wd", "rescale_grad", "clip_gradient"]
type instance OptimizerCst SGD_Mom_Opt dt args =
HasArgs (OptimizerSym SGD_Mom_Opt) (NDArray dt) args '["momentum", "wd", "rescale_grad", "clip_gradient", "lazy_update"]
instance Optimizer SGD_Mom_Opt where
data OptimizerTag SGD_Mom_Opt = SGD'Mom
makeOptimizer SGD'Mom sch args = do
empty <- newIORef M.empty
return $ SGD_Mom_Opt sch args empty
optimize (SGD_Mom_Opt sch args emaref) symbol weight gradient = do
nup <- use $ untag . mod_statistics . stat_num_upd
let lr = getLR sch nup
untag . mod_statistics . stat_last_lr .= lr
liftIO $ do
ema <- readIORef emaref
momentum <- case M.lookup symbol ema of
Nothing -> do
mom <- prim T._zeros_like (#data := weight .& Nil)
writeIORef emaref (M.insert symbol mom ema)
return mom
Just a -> return a
-- let norm x = prim T._norm (#data := x .& #ord := 1 .& Nil)
-- [w0] <- toVector =<< norm weight
-- [m0] <- toVector =<< norm momentum
-- [g0] <- toVector =<< norm gradient
void $ T._sgd_mom_update (#weight := weight
.& #grad := gradient
.& #mom := momentum
.& #lr := lr .& args)
(Just [weight])
-- [w1] <- toVector =<< norm weight
-- [m1] <- toVector =<< norm momentum
-- traceShowM ("opt", symbol, w0, m0, g0, w1, m1)
-- | ADAM optmizer
data ADAM_Opt dtype where
ADAM_Opt :: (LrScheduler sch, OptimizerCst ADAM_Opt dtype args)
=> sch -> ArgsHMap (OptimizerSym ADAM_Opt) (NDArray dtype) args
-> IORef (M.HashMap Text (NDArray dtype, NDArray dtype))
-> ADAM_Opt dtype
type instance OptimizerSym ADAM_Opt = "_adam_update"
-- 1.0.0 type instance OptimizerCst ADAM_Opt dt args = HasArgs (OptimizerSym ADAM_Opt) args '["beta1", "beta2", "epsilon", "wd", "rescale_grad", "clip_gradient"]
type instance OptimizerCst ADAM_Opt dt args =
HasArgs (OptimizerSym ADAM_Opt) (NDArray dt) args '["beta1", "beta2", "epsilon", "wd", "rescale_grad", "clip_gradient", "lazy_update"]
instance Optimizer ADAM_Opt where
data OptimizerTag ADAM_Opt = ADAM
makeOptimizer ADAM sch args = do
empty <- newIORef M.empty
return $ ADAM_Opt sch args empty
optimize (ADAM_Opt sch args emaref) symbol weight gradient = do
nup <- use $ untag . mod_statistics . stat_num_upd
let lr = getLR sch nup
untag . mod_statistics . stat_last_lr .= lr
liftIO $ do
ema <- readIORef emaref
(moving_avg, moving_var) <- case M.lookup symbol ema of
Nothing -> do
avg <- prim T._zeros_like (#data := weight .& Nil)
var <- prim T._ones_like (#data := weight .& Nil)
writeIORef emaref (M.insert symbol (avg, var) ema)
return (avg, var)
Just (a, v) -> return (a, v)
void $ T._adam_update (#weight := weight
.& #grad := gradient
.& #mean := moving_avg
.& #var := moving_var
.& #lr := lr .& args)
(Just [weight])
#ifdef MXNET_GEQ_10600
data ADAMW_Opt dtype where
ADAMW_Opt :: (LrScheduler sch, OptimizerCst ADAMW_Opt dtype args)
=> sch -> ArgsHMap (OptimizerSym ADAMW_Opt) (NDArray dtype) args
-> IORef (M.HashMap Text (NDArray dtype, NDArray dtype))
-> ADAMW_Opt dtype
type instance OptimizerSym ADAMW_Opt = "__adamw_update"
type instance OptimizerCst ADAMW_Opt dt args =
HasArgs (OptimizerSym ADAMW_Opt) (NDArray dt) args
'["beta1", "beta2", "epsilon", "wd", "eta", "clip_gradient", "rescale_grad"]
instance Optimizer ADAMW_Opt where
data OptimizerTag ADAMW_Opt = ADAMW
makeOptimizer ADAMW sch args = do
empty <- newIORef M.empty
return $ ADAMW_Opt sch args empty
optimize (ADAMW_Opt sch args emaref) symbol weight gradient = do
nup <- use $ untag . mod_statistics . stat_num_upd
let lr = getLR sch nup
untag . mod_statistics . stat_last_lr .= lr
liftIO $ do
ema <- readIORef emaref
(moving_avg, moving_var) <- case M.lookup symbol ema of
Nothing -> do
avg <- prim T._zeros_like (#data := weight .& Nil)
var <- prim T._ones_like (#data := weight .& Nil)
writeIORef emaref (M.insert symbol (avg, var) ema)
return (avg, var)
Just (a, v) -> return (a, v)
void $ T.__adamw_update
(#weight := weight .&
#grad := gradient .&
#mean := moving_avg .&
#var := moving_var .&
#lr := lr .& args)
(Just [weight])
#endif