{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

module Torch.Optim.CppOptim where

import Data.Default.Class
import Foreign.ForeignPtr
import System.Mem (performGC)
import Torch.Autograd
import Torch.Internal.Cast
import Torch.Internal.Class (Castable (..), CppObject (..), CppTuple2 (..), CppTuple3 (..), CppTuple4 (..))
import Torch.Internal.GC (mallocTrim)
import qualified Torch.Internal.Managed.Optim as LibTorch
import qualified Torch.Internal.Type as ATen
import Torch.NN
import qualified Torch.Optim as Optim
import Torch.Tensor

type CppOptimizerRef = ForeignPtr ATen.Optimizer

data CppOptimizerState option = CppOptimizerState option CppOptimizerRef

-- class Optimizer option where
--   initOptimizer :: Parameterized model => option -> model -> IO (OptimizerState option model)
--   step :: Parameterized model => OptimizerState option model -> (model -> IO Tensor) -> IO Tensor
--   -- Returned d depends on the state of optimizer.
--   -- Do not call step function after this function is called.
--   getParams :: Parameterized model => OptimizerState option model -> IO model
--   step (OptimizerState _ optimizer initParams) loss = cast0 (LibTorch.step optimizer trans)
--     where
--       trans :: ForeignPtr ATen.TensorList -> IO (ForeignPtr ATen.Tensor)
--       trans inputs =
--         uncast inputs $ \inputs' -> do
--           (Unsafe ret) <- loss $ replaceParameters initParams $  map (IndependentTensor . Unsafe) inputs'
--           cast ret return
--   getParams (OptimizerState _ optimizer initParams) = fmap (replaceParameters initParams . map (IndependentTensor . Unsafe)) $ cast0 (LibTorch.getParams optimizer)

stepWithGenerator ::
  CppOptimizerState option ->
  ForeignPtr ATen.Generator ->
  ([Tensor] -> ForeignPtr ATen.Generator -> IO (Tensor, ForeignPtr ATen.Generator)) ->
  IO (Tensor, ForeignPtr ATen.Generator)
stepWithGenerator :: forall option.
CppOptimizerState option
-> ForeignPtr Generator
-> ([Tensor]
    -> ForeignPtr Generator -> IO (Tensor, ForeignPtr Generator))
-> IO (Tensor, ForeignPtr Generator)
stepWithGenerator o :: CppOptimizerState option
o@(CppOptimizerState option
_ CppOptimizerRef
ref) ForeignPtr Generator
generator [Tensor]
-> ForeignPtr Generator -> IO (Tensor, ForeignPtr Generator)
loss = do
  (Tensor
v, ForeignPtr Generator
nextGenerator) <- (CppOptimizerRef
 -> ForeignPtr Generator
 -> (ForeignPtr TensorList
     -> ForeignPtr Generator
     -> IO (ForeignPtr (StdTuple '(Tensor, Generator))))
 -> IO (ForeignPtr (StdTuple '(Tensor, Generator))))
-> CppOptimizerRef
-> ForeignPtr Generator
-> (ForeignPtr TensorList
    -> ForeignPtr Generator
    -> IO (ForeignPtr (StdTuple '(Tensor, Generator))))
-> IO (Tensor, ForeignPtr Generator)
forall a ca x1 cx1 x2 cx2 y cy.
(Castable a ca, Castable x1 cx1, Castable x2 cx2, Castable y cy) =>
(ca -> cx1 -> cx2 -> IO cy) -> a -> x1 -> x2 -> IO y
cast3 CppOptimizerRef
-> ForeignPtr Generator
-> (ForeignPtr TensorList
    -> ForeignPtr Generator
    -> IO (ForeignPtr (StdTuple '(Tensor, Generator))))
-> IO (ForeignPtr (StdTuple '(Tensor, Generator)))
LibTorch.stepWithGenerator CppOptimizerRef
ref ForeignPtr Generator
generator ForeignPtr TensorList
-> ForeignPtr Generator
-> IO (ForeignPtr (StdTuple '(Tensor, Generator)))
loss'
  (Tensor, ForeignPtr Generator) -> IO (Tensor, ForeignPtr Generator)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Tensor
v, ForeignPtr Generator
nextGenerator)
  where
    loss' :: ForeignPtr ATen.TensorList -> ForeignPtr ATen.Generator -> IO (ForeignPtr (ATen.StdTuple '(ATen.Tensor, ATen.Generator)))
    loss' :: ForeignPtr TensorList
-> ForeignPtr Generator
-> IO (ForeignPtr (StdTuple '(Tensor, Generator)))
loss' ForeignPtr TensorList
params ForeignPtr Generator
gen = do
      (Tensor
v :: Tensor, ForeignPtr Generator
gen') <- ForeignPtr TensorList
-> ([Tensor] -> IO (Tensor, ForeignPtr Generator))
-> IO (Tensor, ForeignPtr Generator)
forall r. ForeignPtr TensorList -> ([Tensor] -> IO r) -> IO r
forall a b r. Castable a b => b -> (a -> IO r) -> IO r
uncast ForeignPtr TensorList
params (([Tensor] -> IO (Tensor, ForeignPtr Generator))
 -> IO (Tensor, ForeignPtr Generator))
-> ([Tensor] -> IO (Tensor, ForeignPtr Generator))
-> IO (Tensor, ForeignPtr Generator)
forall a b. (a -> b) -> a -> b
$ \[Tensor]
params' -> [Tensor]
-> ForeignPtr Generator -> IO (Tensor, ForeignPtr Generator)
loss [Tensor]
params' ForeignPtr Generator
gen
      ForeignPtr Tensor
v' <- Tensor
-> (ForeignPtr Tensor -> IO (ForeignPtr Tensor))
-> IO (ForeignPtr Tensor)
forall r. Tensor -> (ForeignPtr Tensor -> IO r) -> IO r
forall a b r. Castable a b => a -> (b -> IO r) -> IO r
cast Tensor
v ForeignPtr Tensor -> IO (ForeignPtr Tensor)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure :: IO (ForeignPtr ATen.Tensor)
      (ForeignPtr Tensor, ForeignPtr Generator)
-> (ForeignPtr (StdTuple '(Tensor, Generator))
    -> IO (ForeignPtr (StdTuple '(Tensor, Generator))))
-> IO (ForeignPtr (StdTuple '(Tensor, Generator)))
forall r.
(ForeignPtr Tensor, ForeignPtr Generator)
-> (ForeignPtr (StdTuple '(Tensor, Generator)) -> IO r) -> IO r
forall a b r. Castable a b => a -> (b -> IO r) -> IO r
cast (ForeignPtr Tensor
v', ForeignPtr Generator
gen') ForeignPtr (StdTuple '(Tensor, Generator))
-> IO (ForeignPtr (StdTuple '(Tensor, Generator)))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

class CppOptimizer option where
  initOptimizer :: Parameterized model => option -> model -> IO (CppOptimizerState option)
  unsafeStep :: Parameterized model => model -> CppOptimizerState option -> Tensor -> IO (model, CppOptimizerState option)
  unsafeStep model
model o :: CppOptimizerState option
o@(CppOptimizerState option
_ CppOptimizerRef
optimizer) Tensor
loss = do
    [ForeignPtr Tensor]
v <- (CppOptimizerRef
 -> ForeignPtr Tensor -> IO (ForeignPtr TensorList))
-> CppOptimizerRef -> Tensor -> IO [ForeignPtr Tensor]
forall a ca x1 cx1 y cy.
(Castable a ca, Castable x1 cx1, Castable y cy) =>
(ca -> cx1 -> IO cy) -> a -> x1 -> IO y
cast2 CppOptimizerRef -> ForeignPtr Tensor -> IO (ForeignPtr TensorList)
LibTorch.unsafeStep CppOptimizerRef
optimizer Tensor
loss
    let newModel :: model
newModel = model -> [Parameter] -> model
forall f. Parameterized f => f -> [Parameter] -> f
replaceParameters model
model ([Parameter] -> model) -> [Parameter] -> model
forall a b. (a -> b) -> a -> b
$ (ForeignPtr Tensor -> Parameter)
-> [ForeignPtr Tensor] -> [Parameter]
forall a b. (a -> b) -> [a] -> [b]
map (Tensor -> Parameter
IndependentTensor (Tensor -> Parameter)
-> (ForeignPtr Tensor -> Tensor) -> ForeignPtr Tensor -> Parameter
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ForeignPtr Tensor -> Tensor
Unsafe) [ForeignPtr Tensor]
v
    (model, CppOptimizerState option)
-> IO (model, CppOptimizerState option)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (model
newModel, CppOptimizerState option
o)

instance {-# OVERLAPS #-} CppOptimizer option => Optim.Optimizer (CppOptimizerState option) where
  step :: Tensor
-> Gradients
-> [Tensor]
-> CppOptimizerState option
-> ([Tensor], CppOptimizerState option)
step = [Char]
-> Tensor
-> Gradients
-> [Tensor]
-> CppOptimizerState option
-> ([Tensor], CppOptimizerState option)
forall a. HasCallStack => [Char] -> a
error [Char]
"step is not implemented for CppOptimizer."
  runStep :: forall model.
Parameterized model =>
model
-> CppOptimizerState option
-> Tensor
-> Tensor
-> IO (model, CppOptimizerState option)
runStep model
paramState CppOptimizerState option
optState Tensor
lossValue Tensor
lr = do
    IO ()
performGC
    CInt -> IO ()
mallocTrim CInt
0
    model
-> CppOptimizerState option
-> Tensor
-> IO (model, CppOptimizerState option)
forall model.
Parameterized model =>
model
-> CppOptimizerState option
-> Tensor
-> IO (model, CppOptimizerState option)
forall option model.
(CppOptimizer option, Parameterized model) =>
model
-> CppOptimizerState option
-> Tensor
-> IO (model, CppOptimizerState option)
unsafeStep model
paramState CppOptimizerState option
optState Tensor
lossValue

  runStep' :: forall model.
Parameterized model =>
model
-> CppOptimizerState option
-> Gradients
-> Tensor
-> IO (model, CppOptimizerState option)
runStep' = [Char]
-> model
-> CppOptimizerState option
-> Gradients
-> Tensor
-> IO (model, CppOptimizerState option)
forall a. HasCallStack => [Char] -> a
error [Char]
"runStep' is not implemented for CppOptimizer."

data AdagradOptions = AdagradOptions
  { AdagradOptions -> Double
adagradLr :: Double,
    AdagradOptions -> Double
adagradLrDecay :: Double,
    AdagradOptions -> Double
adagradWeightDecay :: Double,
    AdagradOptions -> Double
adagradInitialAccumulatorValue :: Double,
    AdagradOptions -> Double
adagradEps :: Double
  }
  deriving (Int -> AdagradOptions -> ShowS
[AdagradOptions] -> ShowS
AdagradOptions -> [Char]
(Int -> AdagradOptions -> ShowS)
-> (AdagradOptions -> [Char])
-> ([AdagradOptions] -> ShowS)
-> Show AdagradOptions
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> AdagradOptions -> ShowS
showsPrec :: Int -> AdagradOptions -> ShowS
$cshow :: AdagradOptions -> [Char]
show :: AdagradOptions -> [Char]
$cshowList :: [AdagradOptions] -> ShowS
showList :: [AdagradOptions] -> ShowS
Show, AdagradOptions -> AdagradOptions -> Bool
(AdagradOptions -> AdagradOptions -> Bool)
-> (AdagradOptions -> AdagradOptions -> Bool) -> Eq AdagradOptions
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: AdagradOptions -> AdagradOptions -> Bool
== :: AdagradOptions -> AdagradOptions -> Bool
$c/= :: AdagradOptions -> AdagradOptions -> Bool
/= :: AdagradOptions -> AdagradOptions -> Bool
Eq)

instance Default AdagradOptions where
  def :: AdagradOptions
def =
    AdagradOptions
      { adagradLr :: Double
adagradLr = Double
1e-2,
        adagradLrDecay :: Double
adagradLrDecay = Double
0,
        adagradWeightDecay :: Double
adagradWeightDecay = Double
0,
        adagradInitialAccumulatorValue :: Double
adagradInitialAccumulatorValue = Double
0,
        adagradEps :: Double
adagradEps = Double
1e-10
      }

instance CppOptimizer AdagradOptions where
  initOptimizer :: forall model.
Parameterized model =>
AdagradOptions -> model -> IO (CppOptimizerState AdagradOptions)
initOptimizer opt :: AdagradOptions
opt@AdagradOptions {Double
adagradLr :: AdagradOptions -> Double
adagradLrDecay :: AdagradOptions -> Double
adagradWeightDecay :: AdagradOptions -> Double
adagradInitialAccumulatorValue :: AdagradOptions -> Double
adagradEps :: AdagradOptions -> Double
adagradLr :: Double
adagradLrDecay :: Double
adagradWeightDecay :: Double
adagradInitialAccumulatorValue :: Double
adagradEps :: Double
..} model
initParams = do
    CppOptimizerRef
v <- (CDouble
 -> CDouble
 -> CDouble
 -> CDouble
 -> CDouble
 -> ForeignPtr TensorList
 -> IO CppOptimizerRef)
-> Double
-> Double
-> Double
-> Double
-> Double
-> [Tensor]
-> IO CppOptimizerRef
forall a ca x1 cx1 x2 cx2 x3 cx3 x4 cx4 x5 cx5 y cy.
(Castable a ca, Castable x1 cx1, Castable x2 cx2, Castable x3 cx3,
 Castable x4 cx4, Castable x5 cx5, Castable y cy) =>
(ca -> cx1 -> cx2 -> cx3 -> cx4 -> cx5 -> IO cy)
-> a -> x1 -> x2 -> x3 -> x4 -> x5 -> IO y
cast6 CDouble
-> CDouble
-> CDouble
-> CDouble
-> CDouble
-> ForeignPtr TensorList
-> IO CppOptimizerRef
LibTorch.adagrad Double
adagradLr Double
adagradLrDecay Double
adagradWeightDecay Double
adagradInitialAccumulatorValue Double
adagradEps [Tensor]
initParams'
    CppOptimizerState AdagradOptions
-> IO (CppOptimizerState AdagradOptions)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (CppOptimizerState AdagradOptions
 -> IO (CppOptimizerState AdagradOptions))
-> CppOptimizerState AdagradOptions
-> IO (CppOptimizerState AdagradOptions)
forall a b. (a -> b) -> a -> b
$ AdagradOptions
-> CppOptimizerRef -> CppOptimizerState AdagradOptions
forall option.
option -> CppOptimizerRef -> CppOptimizerState option
CppOptimizerState AdagradOptions
opt CppOptimizerRef
v
    where
      initParams' :: [Tensor]
initParams' = (Parameter -> Tensor) -> [Parameter] -> [Tensor]
forall a b. (a -> b) -> [a] -> [b]
map Parameter -> Tensor
toDependent ([Parameter] -> [Tensor]) -> [Parameter] -> [Tensor]
forall a b. (a -> b) -> a -> b
$ model -> [Parameter]
forall f. Parameterized f => f -> [Parameter]
flattenParameters model
initParams

data AdamOptions = AdamOptions
  { AdamOptions -> Double
adamLr :: Double,
    AdamOptions -> (Double, Double)
adamBetas :: (Double, Double),
    AdamOptions -> Double
adamEps :: Double,
    AdamOptions -> Double
adamWeightDecay :: Double,
    AdamOptions -> Bool
adamAmsgrad :: Bool
  }
  deriving (Int -> AdamOptions -> ShowS
[AdamOptions] -> ShowS
AdamOptions -> [Char]
(Int -> AdamOptions -> ShowS)
-> (AdamOptions -> [Char])
-> ([AdamOptions] -> ShowS)
-> Show AdamOptions
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> AdamOptions -> ShowS
showsPrec :: Int -> AdamOptions -> ShowS
$cshow :: AdamOptions -> [Char]
show :: AdamOptions -> [Char]
$cshowList :: [AdamOptions] -> ShowS
showList :: [AdamOptions] -> ShowS
Show, AdamOptions -> AdamOptions -> Bool
(AdamOptions -> AdamOptions -> Bool)
-> (AdamOptions -> AdamOptions -> Bool) -> Eq AdamOptions
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: AdamOptions -> AdamOptions -> Bool
== :: AdamOptions -> AdamOptions -> Bool
$c/= :: AdamOptions -> AdamOptions -> Bool
/= :: AdamOptions -> AdamOptions -> Bool
Eq)

instance Default AdamOptions where
  def :: AdamOptions
def =
    AdamOptions
      { adamLr :: Double
adamLr = Double
1e-3,
        adamBetas :: (Double, Double)
adamBetas = (Double
0.9, Double
0.999),
        adamEps :: Double
adamEps = Double
1e-8,
        adamWeightDecay :: Double
adamWeightDecay = Double
0,
        adamAmsgrad :: Bool
adamAmsgrad = Bool
False
      }

instance CppOptimizer AdamOptions where
  initOptimizer :: forall model.
Parameterized model =>
AdamOptions -> model -> IO (CppOptimizerState AdamOptions)
initOptimizer opt :: AdamOptions
opt@AdamOptions {Bool
Double
(Double, Double)
adamLr :: AdamOptions -> Double
adamBetas :: AdamOptions -> (Double, Double)
adamEps :: AdamOptions -> Double
adamWeightDecay :: AdamOptions -> Double
adamAmsgrad :: AdamOptions -> Bool
adamLr :: Double
adamBetas :: (Double, Double)
adamEps :: Double
adamWeightDecay :: Double
adamAmsgrad :: Bool
..} model
initParams = do
    CppOptimizerRef
v <- (CDouble
 -> CDouble
 -> CDouble
 -> CDouble
 -> CDouble
 -> CBool
 -> ForeignPtr TensorList
 -> IO CppOptimizerRef)
-> Double
-> Double
-> Double
-> Double
-> Double
-> Bool
-> [Tensor]
-> IO CppOptimizerRef
forall a ca x1 cx1 x2 cx2 x3 cx3 x4 cx4 x5 cx5 x6 cx6 y cy.
(Castable a ca, Castable x1 cx1, Castable x2 cx2, Castable x3 cx3,
 Castable x4 cx4, Castable x5 cx5, Castable x6 cx6,
 Castable y cy) =>
(ca -> cx1 -> cx2 -> cx3 -> cx4 -> cx5 -> cx6 -> IO cy)
-> a -> x1 -> x2 -> x3 -> x4 -> x5 -> x6 -> IO y
cast7 CDouble
-> CDouble
-> CDouble
-> CDouble
-> CDouble
-> CBool
-> ForeignPtr TensorList
-> IO CppOptimizerRef
LibTorch.adam Double
adamLr ((Double, Double) -> Double
forall a b. (a, b) -> a
fst (Double, Double)
adamBetas) ((Double, Double) -> Double
forall a b. (a, b) -> b
snd (Double, Double)
adamBetas) Double
adamEps Double
adamWeightDecay Bool
adamAmsgrad [Tensor]
initParams'
    CppOptimizerState AdamOptions -> IO (CppOptimizerState AdamOptions)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (CppOptimizerState AdamOptions
 -> IO (CppOptimizerState AdamOptions))
-> CppOptimizerState AdamOptions
-> IO (CppOptimizerState AdamOptions)
forall a b. (a -> b) -> a -> b
$ AdamOptions -> CppOptimizerRef -> CppOptimizerState AdamOptions
forall option.
option -> CppOptimizerRef -> CppOptimizerState option
CppOptimizerState AdamOptions
opt CppOptimizerRef
v
    where
      initParams' :: [Tensor]
initParams' = (Parameter -> Tensor) -> [Parameter] -> [Tensor]
forall a b. (a -> b) -> [a] -> [b]
map Parameter -> Tensor
toDependent ([Parameter] -> [Tensor]) -> [Parameter] -> [Tensor]
forall a b. (a -> b) -> a -> b
$ model -> [Parameter]
forall f. Parameterized f => f -> [Parameter]
flattenParameters model
initParams

data AdamwOptions = AdamwOptions
  { AdamwOptions -> Double
adamwLr :: Double,
    AdamwOptions -> (Double, Double)
adamwBetas :: (Double, Double),
    AdamwOptions -> Double
adamwEps :: Double,
    AdamwOptions -> Double
adamwWeightDecay :: Double,
    AdamwOptions -> Bool
adamwAmsgrad :: Bool
  }
  deriving (Int -> AdamwOptions -> ShowS
[AdamwOptions] -> ShowS
AdamwOptions -> [Char]
(Int -> AdamwOptions -> ShowS)
-> (AdamwOptions -> [Char])
-> ([AdamwOptions] -> ShowS)
-> Show AdamwOptions
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> AdamwOptions -> ShowS
showsPrec :: Int -> AdamwOptions -> ShowS
$cshow :: AdamwOptions -> [Char]
show :: AdamwOptions -> [Char]
$cshowList :: [AdamwOptions] -> ShowS
showList :: [AdamwOptions] -> ShowS
Show, AdamwOptions -> AdamwOptions -> Bool
(AdamwOptions -> AdamwOptions -> Bool)
-> (AdamwOptions -> AdamwOptions -> Bool) -> Eq AdamwOptions
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: AdamwOptions -> AdamwOptions -> Bool
== :: AdamwOptions -> AdamwOptions -> Bool
$c/= :: AdamwOptions -> AdamwOptions -> Bool
/= :: AdamwOptions -> AdamwOptions -> Bool
Eq)

instance Default AdamwOptions where
  def :: AdamwOptions
def =
    AdamwOptions
      { adamwLr :: Double
adamwLr = Double
1e-3,
        adamwBetas :: (Double, Double)
adamwBetas = (Double
0.9, Double
0.999),
        adamwEps :: Double
adamwEps = Double
1e-8,
        adamwWeightDecay :: Double
adamwWeightDecay = Double
1e-2,
        adamwAmsgrad :: Bool
adamwAmsgrad = Bool
False
      }

instance CppOptimizer AdamwOptions where
  initOptimizer :: forall model.
Parameterized model =>
AdamwOptions -> model -> IO (CppOptimizerState AdamwOptions)
initOptimizer opt :: AdamwOptions
opt@AdamwOptions {Bool
Double
(Double, Double)
adamwLr :: AdamwOptions -> Double
adamwBetas :: AdamwOptions -> (Double, Double)
adamwEps :: AdamwOptions -> Double
adamwWeightDecay :: AdamwOptions -> Double
adamwAmsgrad :: AdamwOptions -> Bool
adamwLr :: Double
adamwBetas :: (Double, Double)
adamwEps :: Double
adamwWeightDecay :: Double
adamwAmsgrad :: Bool
..} model
initParams = do
    CppOptimizerRef
v <- (CDouble
 -> CDouble
 -> CDouble
 -> CDouble
 -> CDouble
 -> CBool
 -> ForeignPtr TensorList
 -> IO CppOptimizerRef)
-> Double
-> Double
-> Double
-> Double
-> Double
-> Bool
-> [Tensor]
-> IO CppOptimizerRef
forall a ca x1 cx1 x2 cx2 x3 cx3 x4 cx4 x5 cx5 x6 cx6 y cy.
(Castable a ca, Castable x1 cx1, Castable x2 cx2, Castable x3 cx3,
 Castable x4 cx4, Castable x5 cx5, Castable x6 cx6,
 Castable y cy) =>
(ca -> cx1 -> cx2 -> cx3 -> cx4 -> cx5 -> cx6 -> IO cy)
-> a -> x1 -> x2 -> x3 -> x4 -> x5 -> x6 -> IO y
cast7 CDouble
-> CDouble
-> CDouble
-> CDouble
-> CDouble
-> CBool
-> ForeignPtr TensorList
-> IO CppOptimizerRef
LibTorch.adamw Double
adamwLr ((Double, Double) -> Double
forall a b. (a, b) -> a
fst (Double, Double)
adamwBetas) ((Double, Double) -> Double
forall a b. (a, b) -> b
snd (Double, Double)
adamwBetas) Double
adamwEps Double
adamwWeightDecay Bool
adamwAmsgrad [Tensor]
initParams'
    CppOptimizerState AdamwOptions
-> IO (CppOptimizerState AdamwOptions)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (CppOptimizerState AdamwOptions
 -> IO (CppOptimizerState AdamwOptions))
-> CppOptimizerState AdamwOptions
-> IO (CppOptimizerState AdamwOptions)
forall a b. (a -> b) -> a -> b
$ AdamwOptions -> CppOptimizerRef -> CppOptimizerState AdamwOptions
forall option.
option -> CppOptimizerRef -> CppOptimizerState option
CppOptimizerState AdamwOptions
opt CppOptimizerRef
v
    where
      initParams' :: [Tensor]
initParams' = (Parameter -> Tensor) -> [Parameter] -> [Tensor]
forall a b. (a -> b) -> [a] -> [b]
map Parameter -> Tensor
toDependent ([Parameter] -> [Tensor]) -> [Parameter] -> [Tensor]
forall a b. (a -> b) -> a -> b
$ model -> [Parameter]
forall f. Parameterized f => f -> [Parameter]
flattenParameters model
initParams

data LbfgsOptions = LbfgsOptions
  { LbfgsOptions -> Double
lbfgsLr :: Double,
    LbfgsOptions -> Int
lbfgsMaxIter :: Int,
    LbfgsOptions -> Int
lbfgsMaxEval :: Int,
    LbfgsOptions -> Double
lbfgsToleranceGrad :: Double,
    LbfgsOptions -> Double
lbfgsToleranceChange :: Double,
    LbfgsOptions -> Int
lbfgsHistorySize :: Int,
    LbfgsOptions -> Maybe [Char]
lbfgsLineSearchFn :: Maybe String
  }
  deriving (Int -> LbfgsOptions -> ShowS
[LbfgsOptions] -> ShowS
LbfgsOptions -> [Char]
(Int -> LbfgsOptions -> ShowS)
-> (LbfgsOptions -> [Char])
-> ([LbfgsOptions] -> ShowS)
-> Show LbfgsOptions
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> LbfgsOptions -> ShowS
showsPrec :: Int -> LbfgsOptions -> ShowS
$cshow :: LbfgsOptions -> [Char]
show :: LbfgsOptions -> [Char]
$cshowList :: [LbfgsOptions] -> ShowS
showList :: [LbfgsOptions] -> ShowS
Show, LbfgsOptions -> LbfgsOptions -> Bool
(LbfgsOptions -> LbfgsOptions -> Bool)
-> (LbfgsOptions -> LbfgsOptions -> Bool) -> Eq LbfgsOptions
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: LbfgsOptions -> LbfgsOptions -> Bool
== :: LbfgsOptions -> LbfgsOptions -> Bool
$c/= :: LbfgsOptions -> LbfgsOptions -> Bool
/= :: LbfgsOptions -> LbfgsOptions -> Bool
Eq)

instance Default LbfgsOptions where
  def :: LbfgsOptions
def =
    LbfgsOptions
      { lbfgsLr :: Double
lbfgsLr = Double
1,
        lbfgsMaxIter :: Int
lbfgsMaxIter = Int
20,
        lbfgsMaxEval :: Int
lbfgsMaxEval = (Int
20 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
5) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
4,
        lbfgsToleranceGrad :: Double
lbfgsToleranceGrad = Double
1e-7,
        lbfgsToleranceChange :: Double
lbfgsToleranceChange = Double
1e-9,
        lbfgsHistorySize :: Int
lbfgsHistorySize = Int
100,
        lbfgsLineSearchFn :: Maybe [Char]
lbfgsLineSearchFn = Maybe [Char]
forall a. Maybe a
Nothing
      }

instance CppOptimizer LbfgsOptions where
  initOptimizer :: forall model.
Parameterized model =>
LbfgsOptions -> model -> IO (CppOptimizerState LbfgsOptions)
initOptimizer opt :: LbfgsOptions
opt@LbfgsOptions {Double
Int
Maybe [Char]
lbfgsLr :: LbfgsOptions -> Double
lbfgsMaxIter :: LbfgsOptions -> Int
lbfgsMaxEval :: LbfgsOptions -> Int
lbfgsToleranceGrad :: LbfgsOptions -> Double
lbfgsToleranceChange :: LbfgsOptions -> Double
lbfgsHistorySize :: LbfgsOptions -> Int
lbfgsLineSearchFn :: LbfgsOptions -> Maybe [Char]
lbfgsLr :: Double
lbfgsMaxIter :: Int
lbfgsMaxEval :: Int
lbfgsToleranceGrad :: Double
lbfgsToleranceChange :: Double
lbfgsHistorySize :: Int
lbfgsLineSearchFn :: Maybe [Char]
..} model
initParams = do
    CppOptimizerRef
v <- (CDouble
 -> CInt
 -> CInt
 -> CDouble
 -> CDouble
 -> CInt
 -> Maybe (ForeignPtr StdString)
 -> ForeignPtr TensorList
 -> IO CppOptimizerRef)
-> Double
-> Int
-> Int
-> Double
-> Double
-> Int
-> Maybe [Char]
-> [Tensor]
-> IO CppOptimizerRef
forall a ca x1 cx1 x2 cx2 x3 cx3 x4 cx4 x5 cx5 x6 cx6 x7 cx7 y cy.
(Castable a ca, Castable x1 cx1, Castable x2 cx2, Castable x3 cx3,
 Castable x4 cx4, Castable x5 cx5, Castable x6 cx6, Castable x7 cx7,
 Castable y cy) =>
(ca -> cx1 -> cx2 -> cx3 -> cx4 -> cx5 -> cx6 -> cx7 -> IO cy)
-> a -> x1 -> x2 -> x3 -> x4 -> x5 -> x6 -> x7 -> IO y
cast8 CDouble
-> CInt
-> CInt
-> CDouble
-> CDouble
-> CInt
-> Maybe (ForeignPtr StdString)
-> ForeignPtr TensorList
-> IO CppOptimizerRef
LibTorch.lbfgs Double
lbfgsLr Int
lbfgsMaxIter Int
lbfgsMaxEval Double
lbfgsToleranceGrad Double
lbfgsToleranceChange Int
lbfgsHistorySize Maybe [Char]
lbfgsLineSearchFn [Tensor]
initParams'
    CppOptimizerState LbfgsOptions
-> IO (CppOptimizerState LbfgsOptions)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (CppOptimizerState LbfgsOptions
 -> IO (CppOptimizerState LbfgsOptions))
-> CppOptimizerState LbfgsOptions
-> IO (CppOptimizerState LbfgsOptions)
forall a b. (a -> b) -> a -> b
$ LbfgsOptions -> CppOptimizerRef -> CppOptimizerState LbfgsOptions
forall option.
option -> CppOptimizerRef -> CppOptimizerState option
CppOptimizerState LbfgsOptions
opt CppOptimizerRef
v
    where
      initParams' :: [Tensor]
initParams' = (Parameter -> Tensor) -> [Parameter] -> [Tensor]
forall a b. (a -> b) -> [a] -> [b]
map Parameter -> Tensor
toDependent ([Parameter] -> [Tensor]) -> [Parameter] -> [Tensor]
forall a b. (a -> b) -> a -> b
$ model -> [Parameter]
forall f. Parameterized f => f -> [Parameter]
flattenParameters model
initParams

data RmspropOptions = RmspropOptions
  { RmspropOptions -> Double
rmspropLr :: Double,
    RmspropOptions -> Double
rmspropAlpha :: Double,
    RmspropOptions -> Double
rmspropEps :: Double,
    RmspropOptions -> Double
rmspropWeightDecay :: Double,
    RmspropOptions -> Double
rmspropMomentum :: Double,
    RmspropOptions -> Bool
rmspropCentered :: Bool
  }
  deriving (Int -> RmspropOptions -> ShowS
[RmspropOptions] -> ShowS
RmspropOptions -> [Char]
(Int -> RmspropOptions -> ShowS)
-> (RmspropOptions -> [Char])
-> ([RmspropOptions] -> ShowS)
-> Show RmspropOptions
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> RmspropOptions -> ShowS
showsPrec :: Int -> RmspropOptions -> ShowS
$cshow :: RmspropOptions -> [Char]
show :: RmspropOptions -> [Char]
$cshowList :: [RmspropOptions] -> ShowS
showList :: [RmspropOptions] -> ShowS
Show, RmspropOptions -> RmspropOptions -> Bool
(RmspropOptions -> RmspropOptions -> Bool)
-> (RmspropOptions -> RmspropOptions -> Bool) -> Eq RmspropOptions
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: RmspropOptions -> RmspropOptions -> Bool
== :: RmspropOptions -> RmspropOptions -> Bool
$c/= :: RmspropOptions -> RmspropOptions -> Bool
/= :: RmspropOptions -> RmspropOptions -> Bool
Eq)

instance Default RmspropOptions where
  def :: RmspropOptions
def =
    RmspropOptions
      { rmspropLr :: Double
rmspropLr = Double
1e-2,
        rmspropAlpha :: Double
rmspropAlpha = Double
0.99,
        rmspropEps :: Double
rmspropEps = Double
1e-8,
        rmspropWeightDecay :: Double
rmspropWeightDecay = Double
0,
        rmspropMomentum :: Double
rmspropMomentum = Double
0,
        rmspropCentered :: Bool
rmspropCentered = Bool
False
      }

instance CppOptimizer RmspropOptions where
  initOptimizer :: forall model.
Parameterized model =>
RmspropOptions -> model -> IO (CppOptimizerState RmspropOptions)
initOptimizer opt :: RmspropOptions
opt@RmspropOptions {Bool
Double
rmspropLr :: RmspropOptions -> Double
rmspropAlpha :: RmspropOptions -> Double
rmspropEps :: RmspropOptions -> Double
rmspropWeightDecay :: RmspropOptions -> Double
rmspropMomentum :: RmspropOptions -> Double
rmspropCentered :: RmspropOptions -> Bool
rmspropLr :: Double
rmspropAlpha :: Double
rmspropEps :: Double
rmspropWeightDecay :: Double
rmspropMomentum :: Double
rmspropCentered :: Bool
..} model
initParams = do
    CppOptimizerRef
v <- (CDouble
 -> CDouble
 -> CDouble
 -> CDouble
 -> CDouble
 -> CBool
 -> ForeignPtr TensorList
 -> IO CppOptimizerRef)
-> Double
-> Double
-> Double
-> Double
-> Double
-> Bool
-> [Tensor]
-> IO CppOptimizerRef
forall a ca x1 cx1 x2 cx2 x3 cx3 x4 cx4 x5 cx5 x6 cx6 y cy.
(Castable a ca, Castable x1 cx1, Castable x2 cx2, Castable x3 cx3,
 Castable x4 cx4, Castable x5 cx5, Castable x6 cx6,
 Castable y cy) =>
(ca -> cx1 -> cx2 -> cx3 -> cx4 -> cx5 -> cx6 -> IO cy)
-> a -> x1 -> x2 -> x3 -> x4 -> x5 -> x6 -> IO y
cast7 CDouble
-> CDouble
-> CDouble
-> CDouble
-> CDouble
-> CBool
-> ForeignPtr TensorList
-> IO CppOptimizerRef
LibTorch.rmsprop Double
rmspropLr Double
rmspropAlpha Double
rmspropEps Double
rmspropWeightDecay Double
rmspropMomentum Bool
rmspropCentered [Tensor]
initParams'
    CppOptimizerState RmspropOptions
-> IO (CppOptimizerState RmspropOptions)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (CppOptimizerState RmspropOptions
 -> IO (CppOptimizerState RmspropOptions))
-> CppOptimizerState RmspropOptions
-> IO (CppOptimizerState RmspropOptions)
forall a b. (a -> b) -> a -> b
$ RmspropOptions
-> CppOptimizerRef -> CppOptimizerState RmspropOptions
forall option.
option -> CppOptimizerRef -> CppOptimizerState option
CppOptimizerState RmspropOptions
opt CppOptimizerRef
v
    where
      initParams' :: [Tensor]
initParams' = (Parameter -> Tensor) -> [Parameter] -> [Tensor]
forall a b. (a -> b) -> [a] -> [b]
map Parameter -> Tensor
toDependent ([Parameter] -> [Tensor]) -> [Parameter] -> [Tensor]
forall a b. (a -> b) -> a -> b
$ model -> [Parameter]
forall f. Parameterized f => f -> [Parameter]
flattenParameters model
initParams

data SGDOptions = SGDOptions
  { SGDOptions -> Double
sgdLr :: Double,
    SGDOptions -> Double
sgdMomentum :: Double,
    SGDOptions -> Double
sgdDampening :: Double,
    SGDOptions -> Double
sgdWeightDecay :: Double,
    SGDOptions -> Bool
sgdNesterov :: Bool
  }
  deriving (Int -> SGDOptions -> ShowS
[SGDOptions] -> ShowS
SGDOptions -> [Char]
(Int -> SGDOptions -> ShowS)
-> (SGDOptions -> [Char])
-> ([SGDOptions] -> ShowS)
-> Show SGDOptions
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> SGDOptions -> ShowS
showsPrec :: Int -> SGDOptions -> ShowS
$cshow :: SGDOptions -> [Char]
show :: SGDOptions -> [Char]
$cshowList :: [SGDOptions] -> ShowS
showList :: [SGDOptions] -> ShowS
Show, SGDOptions -> SGDOptions -> Bool
(SGDOptions -> SGDOptions -> Bool)
-> (SGDOptions -> SGDOptions -> Bool) -> Eq SGDOptions
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: SGDOptions -> SGDOptions -> Bool
== :: SGDOptions -> SGDOptions -> Bool
$c/= :: SGDOptions -> SGDOptions -> Bool
/= :: SGDOptions -> SGDOptions -> Bool
Eq)

instance Default SGDOptions where
  def :: SGDOptions
def =
    SGDOptions
      { sgdLr :: Double
sgdLr = Double
1e-3,
        sgdMomentum :: Double
sgdMomentum = Double
0,
        sgdDampening :: Double
sgdDampening = Double
0,
        sgdWeightDecay :: Double
sgdWeightDecay = Double
0,
        sgdNesterov :: Bool
sgdNesterov = Bool
False
      }

instance CppOptimizer SGDOptions where
  initOptimizer :: forall model.
Parameterized model =>
SGDOptions -> model -> IO (CppOptimizerState SGDOptions)
initOptimizer opt :: SGDOptions
opt@SGDOptions {Bool
Double
sgdLr :: SGDOptions -> Double
sgdMomentum :: SGDOptions -> Double
sgdDampening :: SGDOptions -> Double
sgdWeightDecay :: SGDOptions -> Double
sgdNesterov :: SGDOptions -> Bool
sgdLr :: Double
sgdMomentum :: Double
sgdDampening :: Double
sgdWeightDecay :: Double
sgdNesterov :: Bool
..} model
initParams = do
    CppOptimizerRef
v <- (CDouble
 -> CDouble
 -> CDouble
 -> CDouble
 -> CBool
 -> ForeignPtr TensorList
 -> IO CppOptimizerRef)
-> Double
-> Double
-> Double
-> Double
-> Bool
-> [Tensor]
-> IO CppOptimizerRef
forall a ca x1 cx1 x2 cx2 x3 cx3 x4 cx4 x5 cx5 y cy.
(Castable a ca, Castable x1 cx1, Castable x2 cx2, Castable x3 cx3,
 Castable x4 cx4, Castable x5 cx5, Castable y cy) =>
(ca -> cx1 -> cx2 -> cx3 -> cx4 -> cx5 -> IO cy)
-> a -> x1 -> x2 -> x3 -> x4 -> x5 -> IO y
cast6 CDouble
-> CDouble
-> CDouble
-> CDouble
-> CBool
-> ForeignPtr TensorList
-> IO CppOptimizerRef
LibTorch.sgd Double
sgdLr Double
sgdMomentum Double
sgdDampening Double
sgdWeightDecay Bool
sgdNesterov [Tensor]
initParams'
    CppOptimizerState SGDOptions -> IO (CppOptimizerState SGDOptions)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (CppOptimizerState SGDOptions -> IO (CppOptimizerState SGDOptions))
-> CppOptimizerState SGDOptions
-> IO (CppOptimizerState SGDOptions)
forall a b. (a -> b) -> a -> b
$ SGDOptions -> CppOptimizerRef -> CppOptimizerState SGDOptions
forall option.
option -> CppOptimizerRef -> CppOptimizerState option
CppOptimizerState SGDOptions
opt CppOptimizerRef
v
    where
      initParams' :: [Tensor]
initParams' = (Parameter -> Tensor) -> [Parameter] -> [Tensor]
forall a b. (a -> b) -> [a] -> [b]
map Parameter -> Tensor
toDependent ([Parameter] -> [Tensor]) -> [Parameter] -> [Tensor]
forall a b. (a -> b) -> a -> b
$ model -> [Parameter]
forall f. Parameterized f => f -> [Parameter]
flattenParameters model
initParams

saveState :: CppOptimizerState option -> FilePath -> IO ()
saveState :: forall option. CppOptimizerState option -> [Char] -> IO ()
saveState (CppOptimizerState option
_ CppOptimizerRef
optimizer) [Char]
file = (CppOptimizerRef -> ForeignPtr StdString -> IO ())
-> CppOptimizerRef -> [Char] -> IO ()
forall a ca x1 cx1 y cy.
(Castable a ca, Castable x1 cx1, Castable y cy) =>
(ca -> cx1 -> IO cy) -> a -> x1 -> IO y
cast2 CppOptimizerRef -> ForeignPtr StdString -> IO ()
LibTorch.save CppOptimizerRef
optimizer [Char]
file

loadState :: CppOptimizerState option -> FilePath -> IO ()
loadState :: forall option. CppOptimizerState option -> [Char] -> IO ()
loadState (CppOptimizerState option
_ CppOptimizerRef
optimizer) [Char]
file = (CppOptimizerRef -> ForeignPtr StdString -> IO ())
-> CppOptimizerRef -> [Char] -> IO ()
forall a ca x1 cx1 y cy.
(Castable a ca, Castable x1 cx1, Castable y cy) =>
(ca -> cx1 -> IO cy) -> a -> x1 -> IO y
cast2 CppOptimizerRef -> ForeignPtr StdString -> IO ()
LibTorch.load CppOptimizerRef
optimizer [Char]
file