diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,26 @@ +# Haskell stochastic gradient descent library + +Stochastic gradient descent (SGD) is a method for optimizing a global objective +function defined as a sum of smaller, differentiable functions. In each +iteration of SGD the gradient is calculated based on a subset of the training +dataset. In Haskell, this process can be simply represented as a [fold over a +of subsequent dataset +subsets](https://blog.jle.im/entry/purely-functional-typed-models-1.html) +(singleton elements in the extreme). + +However, it can be beneficial to select the subsequent subsets randomly (e.g., +shuffle the entire dataset before each pass). Moreover, the dataset can be +large enough to make it impractical to store it all in memory. Hence, the +`sgd` library adopts a [pipe](http://hackage.haskell.org/package/pipes)-based +interface in which SGD takes the form of a process consuming dataset subsets +(the so-called mini-batches) and producing a stream of output parameter values. + +The `sgd` library implements several SGD variants (SGD with momentum, AdaDelta, +Adam) and handles heterogeneous parameter representations (vectors, maps, custom +records, etc.). It can be used in combination with automatic differentiation +libraries ([ad](http://hackage.haskell.org/package/ad), +[backprop](http://hackage.haskell.org/package/backprop)), which can be used to +automatically calculate the gradient of the objective function. + +Look at [the hackage repository](http://hackage.haskell.org/package/sgd) for a +library documentation. diff --git a/changelog b/changelog new file mode 100644 --- /dev/null +++ b/changelog @@ -0,0 +1,5 @@ +-*-change-log-*- + +0.8.0 Apr 2019 + * Allow step size decay in Adam + * Parallel calculation of the objective function diff --git a/examples/example1.hs b/examples/example1.hs deleted file mode 100644 --- a/examples/example1.hs +++ /dev/null @@ -1,104 +0,0 @@ -{-# LANGUAGE RecordWildCards #-} - -import Control.Applicative ((<$>), (<*>)) -import Control.Monad (replicateM) -import System.IO (hSetBuffering, stdout, BufferMode (NoBuffering)) -import qualified System.Random as R -import qualified Data.Vector as V -import qualified Data.Vector.Unboxed as U -import qualified Numeric.SGD as S - ------------------------------------------------------------------------------- --- Dataset generation ------------------------------------------------------------------------------- - --- | Element of a dataset. -type Elem = [(Int, Double)] - --- | Random dataset element. -elemR - :: Int -- ^ Maximum number of element items - -> (Int, Int) -- ^ Range for item's first component - -> (Double, Double) -- ^ Range for item's second component - -> IO Elem -- ^ Result -elemR nMax xr yr = do - n <- R.randomRIO (0, max 0 nMax) - replicateM n ((,) <$> R.randomRIO xr <*> R.randomRIO yr) - --- | Random dataset. -dataSetR - :: Int -- ^ Dataset size - -> Int -- ^ Number of model parameters - -> Int -- ^ Maximum number of items in data element - -> (Double, Double) -- ^ Range for item's second component - -> IO (V.Vector Elem) -- ^ Result -dataSetR m n k yRan = - V.fromList <$> replicateM m (elemR k (0, n-1) yRan) - ------------------------------------------------------------------------------- --- Objective function and gradient ------------------------------------------------------------------------------- - --- | An objective function. The SGD method can be used when --- the objective function is defined in a form of a sum. -goal :: S.Para -> [Elem] -> Double -goal para = - sum . map perElem - where - perElem xs = sum - [ (para U.! k - x) ^ (2 :: Int) - | (k, x) <- xs ] - --- | Since the goal function has a form of a sum, it is sufficient to define --- the gradient over one element only. The gradient with respect to the dataset --- is a sum of gradients over its individual elements. -grad :: S.Para -> Elem -> S.Grad -grad para xs = S.fromList - -- [ (k, 2 * (x - para U.! k)) - [ (k, 2 * (para U.! k - x)) - | (k, x) <- xs ] - --- | Negate gradient. We use it to find the minimum of the objective function. -negGrad :: (S.Para -> Elem -> S.Grad) - -> (S.Para -> Elem -> S.Grad) -negGrad g para x = fmap negate (g para x) - ------------------------------------------------------------------------------- --- SGD ------------------------------------------------------------------------------- - --- | Notification run by the sgdM function every parameters update. -notify :: S.SgdArgs -> V.Vector Elem -> S.Para -> Int -> IO () -notify S.SgdArgs{..} dataSet para k = - if doneTotal k /= doneTotal (k - 1) - then do - let n = doneTotal k - x = goal para (V.toList dataSet) - putStrLn ("\n" ++ "[" ++ show n ++ "] f = " ++ show x) - else - putStr "." - where - doneTotal :: Int -> Int - doneTotal = floor . done - done :: Int -> Double - done i - = fromIntegral (i * batchSize) - / fromIntegral (V.length dataSet) - --- | Run the monadic version of SGD. -runSgdM - :: Int -- ^ Dataset size - -> Int -- ^ Number of model parameters - -> Int -- ^ Maximum number of items in data element - -> S.SgdArgs -- ^ SGD parameters - -> IO S.Para -runSgdM m n k sgdArgs = do - dataSet <- dataSetR m n k (-10, 10) - let para = U.replicate n 0 - hSetBuffering stdout NoBuffering - S.sgdM sgdArgs (notify sgdArgs dataSet) (negGrad grad) dataSet para - --- | Run the monadic version of SGD with some default parameter values. -main = do - let sgdArgs = S.sgdArgsDefault { S.iterNum = 50 } - runSgdM 1000 1000000 10 sgdArgs diff --git a/sgd.cabal b/sgd.cabal --- a/sgd.cabal +++ b/sgd.cabal @@ -1,52 +1,99 @@ -name: sgd -version: 0.3.7 -synopsis: Stochastic gradient descent -description: - Implementation of a Stochastic Gradient Descent optimization method. - See examples directory in the source package for examples of usage. - . - It is a preliminary implementation of the SGD method and API may change - in future versions. -license: BSD3 -license-file: LICENSE -cabal-version: >= 1.6 -copyright: Copyright (c) 2012 IPI PAN -author: Jakub Waszczuk -maintainer: waszczuk.kuba@gmail.com -stability: experimental -category: Math, Algorithms -homepage: https://github.com/kawu/sgd -build-type: Simple - -extra-source-files: examples/example1.hs +cabal-version: 1.12 -library - hs-source-dirs: src +-- This file has been generated from package.yaml by hpack version 0.31.1. +-- +-- see: https://github.com/sol/hpack +-- +-- hash: 95a0d8b144a4bd5d9864eca97315546b377194d492e5676678cab11540085447 - build-depends: - base >= 4 && < 5 - , containers >= 0.4 && < 0.6 - , vector >= 0.10 && < 0.11 - , random >= 1.0 && < 1.1 - , primitive >= 0.5 && < 0.6 - , logfloat >= 0.12 && < 0.13 - , monad-par >= 0.3.4 && < 0.4 - , deepseq >= 1.3 && < 1.4 - , binary >= 0.5 && < 0.8 - , bytestring >= 0.9 && < 0.11 - , mtl >= 2.0 && < 2.3 - , filepath >= 1.3 && < 1.4 - , temporary >= 1.1 && < 1.2 - , lazy-io >= 0.1 && < 0.2 +name: sgd +version: 0.8.0.3 +synopsis: Stochastic gradient descent library +description: Import "Numeric.SGD" to use the library. +category: Math +homepage: https://github.com/kawu/sgd#readme +bug-reports: https://github.com/kawu/sgd/issues +author: Jakub Waszczuk +maintainer: waszczuk.kuba@gmail.com +copyright: 2012-2019 Jakub Waszczuk +license: BSD3 +license-file: LICENSE +build-type: Simple +extra-source-files: + README.md + changelog - exposed-modules: - Numeric.SGD - , Numeric.SGD.Dataset - , Numeric.SGD.LogSigned - , Numeric.SGD.Grad +source-repository head + type: git + location: https://github.com/kawu/sgd - ghc-options: -Wall -O2 +library + exposed-modules: + Numeric.SGD + Numeric.SGD.AdaDelta + Numeric.SGD.Adam + Numeric.SGD.DataSet + Numeric.SGD.Momentum + Numeric.SGD.ParamSet + Numeric.SGD.Sparse + Numeric.SGD.Sparse.Grad + Numeric.SGD.Sparse.LogSigned + Numeric.SGD.Sparse.Momentum + Numeric.SGD.Type + other-modules: + Paths_sgd + hs-source-dirs: + src + build-depends: + base >=4.7 && <5 + , binary >=0.5 && <0.9 + , bytestring >=0.9 && <0.11 + , containers >=0.4 && <0.7 + , data-default >=0.7 && <0.8 + , deepseq >=1.3 && <1.5 + , filepath >=1.3 && <1.5 + , hmatrix >=0.19 && <0.21 + , logfloat >=0.12 && <0.14 + , monad-par >=0.3.4 && <0.4 + , mtl >=2.0 && <2.3 + , parallel >=3.2 && <3.3 + , pipes >=4.3 && <4.4 + , primitive >=0.5 && <0.7 + , random >=1.0 && <1.2 + , random-shuffle >=0.0.4 && <0.1 + , temporary >=1.1 && <1.4 + , vector >=0.10 && <0.13 + default-language: Haskell2010 -source-repository head - type: git - location: git://github.com/kawu/sgd.git +test-suite vine-test + type: exitcode-stdio-1.0 + main-is: Spec.hs + other-modules: + Paths_sgd + hs-source-dirs: + test + ghc-options: -threaded -rtsopts -with-rtsopts=-N + build-depends: + ad >=4.3 && <4.4 + , base >=4.7 && <5 + , binary >=0.5 && <0.9 + , bytestring >=0.9 && <0.11 + , containers >=0.4 && <0.7 + , data-default >=0.7 && <0.8 + , deepseq >=1.3 && <1.5 + , filepath >=1.3 && <1.5 + , hmatrix >=0.19 && <0.21 + , logfloat >=0.12 && <0.14 + , monad-par >=0.3.4 && <0.4 + , mtl >=2.0 && <2.3 + , parallel >=3.2 && <3.3 + , pipes >=4.3 && <4.4 + , primitive >=0.5 && <0.7 + , random >=1.0 && <1.2 + , random-shuffle >=0.0.4 && <0.1 + , sgd + , tasty >=1.2 && <1.3 + , tasty-hunit >=0.10 && <0.11 + , temporary >=1.1 && <1.4 + , vector >=0.10 && <0.13 + default-language: Haskell2010 diff --git a/src/Numeric/SGD.hs b/src/Numeric/SGD.hs --- a/src/Numeric/SGD.hs +++ b/src/Numeric/SGD.hs @@ -1,136 +1,490 @@ {-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE DeriveGeneric #-} --- | Stochastic gradient descent implementation using mutable --- vectors for efficient update of the parameters vector. --- A user is provided with the immutable vector of parameters --- so he is able to compute the gradient outside of the IO monad. --- Currently only the Gaussian priors are implemented. +-- | Main module of the stochastic gradient descent (SGD) library. -- --- This is a preliminary version of the SGD library and API may change --- in future versions. - +-- SGD is a method for optimizing a global objective function defined as a sum +-- of smaller, differentiable functions. The individual component functions +-- share the same set of parameters, represented by the `ParamSet` class. This +-- allows for heterogeneous parameter representation (vectors, maps, custom +-- records, etc.). +-- +-- The library adopts a `P.Pipe`-based interface in which `SGD` takes the form +-- of a process consuming dataset subsets (the so-called mini-batches) and +-- producing a stream of output parameter values. The library implements +-- different variants of `SGD` (`Mom.momentum`, `Adam.adam`, `Ada.adaDelta`) +-- which can be executed in either the pure context (`run`) or in IO (`runIO`). +-- The use of lower-level pipe-processing combinators (`pipeRan`, `batch`, +-- `result`, etc.) is also possible. +-- +-- To perform SGD, the gradients of the individual functions need to be +-- determined. This can be done manually or automatically, using an automatic +-- differentiation library (, +-- ). +-- module Numeric.SGD -( SgdArgs (..) -, sgdArgsDefault -, Para -, sgd -, module Numeric.SGD.Grad -, module Numeric.SGD.Dataset -) where + ( + -- * Example + -- $example + -- * SGD variants + Mom.momentum + , Ada.adaDelta + , Adam.adam -import Control.Monad (forM_) -import qualified System.Random as R -import qualified Data.Vector.Unboxed as U -import qualified Data.Vector.Unboxed.Mutable as UM -import qualified Control.Monad.Primitive as Prim + -- * Pure SGD + , run -import Numeric.SGD.Grad -import Numeric.SGD.Dataset + -- * IO-based SGD + , Config (..) + , iterNumPerEpoch + , reportObjective + , objectiveWith + , runIO + -- * Combinators + -- ** Input + , pipeSeq + , pipeRan + -- ** Batch + , batch + , batchGradSeq + , batchGradPar + , batchGradPar' + -- ** Output + , result + -- ** Misc + , keepEvery + , decreasingBy --- | SGD parameters controlling the learning process. -data SgdArgs = SgdArgs - { -- | Size of the batch - batchSize :: Int - -- | Regularization variance - , regVar :: Double - -- | Number of iterations - , iterNum :: Double - -- | Initial gain parameter - , gain0 :: Double - -- | After how many iterations over the entire dataset - -- the gain parameter is halved - , tau :: Double } + -- * Re-exports + , def + ) where --- | Default SGD parameter values. -sgdArgsDefault :: SgdArgs -sgdArgsDefault = SgdArgs - { batchSize = 30 - , regVar = 10 - , iterNum = 10 - , gain0 = 1 - , tau = 5 } +import GHC.Generics (Generic) +import GHC.Conc (numCapabilities) +import Numeric.Natural (Natural) --- | Vector of parameters. -type Para = U.Vector Double +-- import qualified System.Random as R +import Control.Monad (when, forever) -- forM_, +import Control.Parallel.Strategies (parMap, rseq, rdeepseq, Strategy) +import Control.DeepSeq (NFData) +import qualified Control.Monad.State.Strict as State --- | Type synonym for mutable vector with Double values. -type MVect = UM.MVector (Prim.PrimState IO) Double +import Data.Functor.Identity (Identity(..)) +import Data.List (foldl1', transpose) --foldl', +-- import qualified Data.IORef as IO +import Data.Default --- | A stochastic gradient descent method. --- A notification function can be used to provide user with --- information about the progress of the learning. -sgd - :: SgdArgs -- ^ SGD parameter values - -> (Para -> Int -> IO ()) -- ^ Notification run every update - -> (Para -> x -> Grad) -- ^ Gradient for dataset element - -> Dataset x -- ^ Dataset - -> Para -- ^ Starting point - -> IO Para -- ^ SGD result -sgd SgdArgs{..} notify mkGrad dataset x0 = do - u <- UM.new (U.length x0) - doIt u 0 (R.mkStdGen 0) =<< U.thaw x0 +import qualified Pipes as P +import qualified Pipes.Prelude as P +import Pipes ((>->)) + +import qualified Numeric.SGD.Momentum as Mom +import qualified Numeric.SGD.AdaDelta as Ada +import qualified Numeric.SGD.Adam as Adam +import Numeric.SGD.Type +import Numeric.SGD.ParamSet +import Numeric.SGD.DataSet + + +{- $example + + Let's say we have a list of functions defined as: + +> funs = [\x -> 0.3*x^2, \x -> -2*x, const 3, sin] + + The global objective (which we want to minimize) is then defined as: + +> objective x = sum $ map ($x) funs + + To perform SGD, we can either manually determine the individual derivatives: + +> derivs = [\x -> 0.6*x, const (-2), const 0, cos] + + or use an automatic differentiation library, for instance: + +> import qualified Numeric.AD as AD +> derivs = +> [ AD.diff $ \x -> 0.3*x^2 +> , AD.diff $ \x -> -2*x +> , AD.diff $ const 3 +> , AD.diff $ sin +> ] + + Finally, `run` allows to approach a (potentially local) minimum of the + global objective function: + +>>> run (momentum def id) (take 10000 $ cycle derivs) 0.0 +4.180177042912455 + + where: + + * @(take 10000 $ cycle derivs)@ is the stream of training examples + * @(momentum def id)@ is the selected SGD variant (`Mom.momentum`), + supplied with the default configuration (`def`) and the function (`id`) + for calculating the gradient from a training example + * @0.0@ is the initial parameter value + +-} + + +------------------------------- +-- Pure SGD +------------------------------- + + +-- | Traverse all the elements in the training data stream in one pass, +-- calculate the subsequent gradients, and apply them progressively starting +-- from the initial parameter values. +-- +-- Consider using `runIO` if your training dataset is large. +run + :: (ParamSet p) + => SGD Identity e p + -- ^ Selected SGD method + -> [e] + -- ^ Training data stream + -> p + -- ^ Initial parameters + -> p +run sgd dataSet p0 = runIdentity $ + result p0 + (P.each dataSet >-> sgd p0) + + +------------------------------- +-- Higher-level SGD +------------------------------- + + +-- | High-level IO-based SGD configuration +data Config = Config + { iterNum :: Natural + -- ^ Number of iteration over the entire training dataset + , batchSize :: Natural + -- ^ Mini-batch size + , batchOverlap :: Natural + -- ^ The number of overlapping elements in subsequent mini-batches + , batchRandom :: Bool + -- ^ Should the mini-batch be selected at random? If not, the subsequent + -- training elements will be picked sequentially. Random selection gives + -- no guarantee of seeing each training sample in every epoch. + , reportEvery :: Double + -- ^ How often the value of the objective function should be reported (with + -- @1@ meaning once per pass over the training data) + } deriving (Show, Eq, Ord, Generic) + +instance Default Config where + def = Config + { iterNum = 100 + , batchSize = 1 + , batchOverlap = 0 + , batchRandom = False + , reportEvery = 1.0 + } + + +-- | Number of new elements in each new batch +batchNew :: Config -> Int +batchNew cfg = max 1 + ( fromIntegral (batchSize cfg) + - fromIntegral (batchOverlap cfg) + ) + + +-- | Calculate the effective number of SGD iterations (and gradient +-- calculations) performed per epoch. +iterNumPerEpoch + :: (Integral a) + => Config + -> a -- ^ Dataset size + -> Double +iterNumPerEpoch cfg size = + fromIntegral size / fromIntegral (batchNew cfg) + + +-- | Report the total objective value on stdout. +reportObjective + :: (ParamSet p) + => (e -> p -> Double) + -- ^ Value of the objective function on a dataset element + -> DataSet e + -- ^ Training dataset + -> p -> IO Double +reportObjective objAt dataSet net = do + q <- objectiveWith objAt dataSet net + putStr $ show q + putStrLn $ " (norm_2 = " ++ show (norm_2 net) ++ ")" + return q + + +-- | Value of the objective function over the entire dataset (i.e. the sum of +-- the objectives on all dataset elements). +objectiveWith + :: (e -> p -> Double) + -- ^ Value of the objective function on a dataset element + -> DataSet e + -- ^ Training dataset + -> p -> IO Double +objectiveWith objAt dataSet net = do + +-- elems <- loadData dataSet +-- let scores = parMap rseq (flip objAt net) elems +-- return $ sum scores + + parts <- partition numCapabilities <$> loadData dataSet + let scores = parMap rseq groupScore parts + return $ sum scores where - -- Gain in k-th iteration. - gain k = (gain0 * tau) / (tau + done k) + groupScore = sum . map (flip objAt net) - -- Number of completed iterations over the full dataset. - done k - = fromIntegral (k * batchSize) - / fromIntegral (size dataset) +-- res <- IO.newIORef 0.0 +-- forM_ [0 .. size dataSet - 1] $ \ix -> do +-- x <- elemAt dataSet ix +-- IO.modifyIORef' res (+ objAt x net) +-- IO.readIORef res - doIt u k stdGen x - | done k > iterNum = do - frozen <- U.unsafeFreeze x - notify frozen k - return frozen + +-- | Perform SGD in the IO monad, regularly reporting the value of the +-- objective function on the entire dataset. A higher-level wrapper which +-- should be convenient to use when the training dataset is large. +-- +-- An alternative is to use the simpler function `run`, or to build a custom +-- SGD pipeline based on lower-level combinators (`pipeSeq`, `batch`, +-- `Adam.adam`, `keepEvery`, `result`, etc.). +runIO + :: (ParamSet p) + => Config + -- ^ SGD configuration + -> SGD IO [e] p + -- ^ SGD pipe consuming mini-batches of dataset elements + -> (p -> IO Double) + -- ^ Quality reporting function (the reporting frequency is specified + -- via `reportEvery`) + -> DataSet e + -- ^ Training dataset + -> p + -- ^ Initial parameter values + -> IO p +runIO cfg@Config{..} sgd reportObj dataSet net0 = do + _ <- reportObj net0 + result net0 $ pipeData dataSet + >-> batch (fromIntegral batchSize) + >-> batchFilter + >-> sgd net0 + >-> keepEvery realReportPeriod + >-> P.take (fromIntegral iterNum) + >-> decreasingBy reportObj + where + -- Data streaming function + pipeData = forever . + if batchRandom + then pipeRan + else pipeSeq + -- Batch stream filter + batchFilter = do + P.await >>= P.yield + keepEvery (batchNew cfg) + -- Iteration (epoch) scaling + realReportPeriod = ceiling $ + reportEvery * iterNumPerEpoch cfg (size dataSet) + + +------------------------------- +-- Lower-level combinators +------------------------------- + + +-- | Pipe all the elements in the dataset sequentially. +pipeSeq :: DataSet e -> P.Producer e IO () +pipeSeq dataSet = do + go (0 :: Int) + where + go k + | k >= size dataSet = return () | otherwise = do - (batch, stdGen') <- sample stdGen batchSize dataset + x <- P.lift $ elemAt dataSet k + P.yield x + go (k+1) - -- Freeze mutable vector of parameters. The frozen version is - -- then supplied to external mkGrad function provided by user. - frozen <- U.unsafeFreeze x - notify frozen k - -- let grad = M.unions (map (mkGrad frozen) batch) - let grad = parUnions (map (mkGrad frozen) batch) - addUp grad u - scale (gain k) u +-- | Pipe all the elements in the dataset in a random order. +pipeRan :: DataSet e -> P.Producer e IO () +pipeRan dataSet0 = do + dataSet <- P.lift $ shuffle dataSet0 + pipeSeq dataSet - x' <- U.unsafeThaw frozen - apply u x' - doIt u (k+1) stdGen' x' +-- | Group dataset elements into (mini-)batches of the given size. +batch :: (Monad m) => Int -> P.Pipe e [e] m () +batch k = flip State.evalStateT [] . forever $ do + x <- P.lift P.await + xs <- State.get + let xs' = take k (x:xs) + when (length xs' == k) $ do + P.lift (P.yield xs') + State.put xs' --- | Add up all gradients and store results in normal domain. -addUp :: Grad -> MVect -> IO () -addUp grad v = do - UM.set v 0 - forM_ (toList grad) $ \(i, x) -> do - y <- UM.unsafeRead v i - UM.unsafeWrite v i (x + y) +-- | Adapt the gradient function to handle (mini-)batches. Relies on the @p@'s +-- `NFData` instance to efficiently calculate gradients in parallel. +batchGradPar + :: (ParamSet p, NFData p) + => (e -> p -> p) + -> ([e] -> p -> p) +batchGradPar = batchGradWith rdeepseq --- | Scale the vector by the given value. -scale :: Double -> MVect -> IO () -scale c v = do - forM_ [0 .. UM.length v - 1] $ \i -> do - y <- UM.unsafeRead v i - UM.unsafeWrite v i (c * y) +-- | A version of `batchGradPar` with no `NFData` constraint. Evaluates the +-- sub-gradients calculated in parallel to weak head normal form. +batchGradPar' + :: (ParamSet p) + => (e -> p -> p) + -> ([e] -> p -> p) +batchGradPar' = batchGradWith rseq --- | Apply gradient to the parameters vector, that is add the first vector to --- the second one. -apply :: MVect -> MVect -> IO () -apply w v = do - forM_ [0 .. UM.length v - 1] $ \i -> do - x <- UM.unsafeRead v i - y <- UM.unsafeRead w i - UM.unsafeWrite v i (x + y) + +-- | Adapt the gradient function to handle (mini-)batches. The sub-gradients +-- of the individual batch elements are evaluated in parallel based on the +-- given `Strategy`. +batchGradWith + :: (ParamSet p) + => Strategy p + -> (e -> p -> p) + -> ([e] -> p -> p) +batchGradWith strategy grad xs param = + case parMap strategy (\e -> grad e param) xs of + [] -> param + -- TODO: the fold is sequential, we could try to parallize it as well. + ps -> foldl1' add ps + + +-- -- | Adapt the gradient function to handle (mini-)batches. The sub-gradients +-- -- of the individual batch elements are evaluated in parallel based on the +-- -- given `Strategy`. +-- batchGradWith +-- :: (ParamSet p) +-- => Strategy p +-- -> (e -> p -> p) +-- -> ([e] -> p -> p) +-- batchGradWith strategy grad xs param = +-- +-- addAll grads +-- +-- where +-- +-- groups = partition numCapabilities xs +-- grads = parMap strategy gradMany groups +-- +-- -- TODO: can we assume here that the group is non-empty? +-- gradMany = foldl1' add . map (\e -> grad e param) +-- +-- addAll [] = param +-- addAll ps = foldl1' add ps + + +-- | Adapt the gradient function to handle (mini-)batches. The function +-- calculates the individual sub-gradients sequentially. +batchGradSeq + :: (ParamSet p) + => (e -> p -> p) + -> ([e] -> p -> p) +batchGradSeq grad xs param = + case map (flip grad param) xs of + [] -> param + ps -> foldl1' add ps + + +-- | Extract the result of the SGD calculation (the last parameter +-- set flowing downstream). +result + :: (Monad m) + => p + -- ^ Default value (in case the stream is empty) + -> P.Producer p m () + -- ^ Stream of parameter sets + -> m p +result pDef = fmap (maybe pDef id) . P.last + + +-- -- | Apply the given monadic function to every @k@-th value flowing downstream. +-- every :: (Monad m) => Int -> (p -> m ()) -> P.Pipe p p m x +-- every k f = do +-- go (1 `mod` k) +-- where +-- go i = do +-- paramSet <- P.await +-- when (i == 0) $ do +-- P.lift $ f paramSet +-- P.yield paramSet +-- go $ (i+1) `mod` k + + +-- | Keep every @k@-th element flowing downstream and discard all the others. +keepEvery :: (Monad m) => Int -> P.Pipe a a m x +keepEvery k = forever $ do + sequence_ $ replicate (k-1) P.await + P.await >>= P.yield +-- keepEvery k = do +-- go (1 `mod` k) +-- where +-- go i = do +-- x <- P.await +-- when (i == 0) $ do +-- P.yield x +-- go $ (i+1) `mod` k + + +-- -- | Keep the elements with the corresponding `True` in the argument list. +-- -- +-- -- TODO: (=) or (==) in the following example? And is this example correct? +-- -- @ +-- -- keep (forever True) = P.id +-- -- @ +-- keep :: (Monad m) => [Bool] -> P.Pipe a a m () +-- keep [] = return () +-- keep (b:bs) = do +-- x <- P.await +-- when b (P.yield x) +-- keep bs +-- +-- +-- -- | Create the mask to `keep` each @k@-th element flowing downstream. +-- every :: Int -> [Bool] +-- every k = cycle $ replicate (k-1) False ++ [True] + + +-- | Make the stream decreasing in the given (monadic) function by discarding +-- elements with values higher than those already seen. +decreasingBy :: (Monad m, Ord a) => (p -> m a) -> P.Pipe p p m x +decreasingBy f = do + x <- P.await + v <- P.lift (f x) + P.yield x + go v + where + go w = do + x <- P.await + v <- P.lift (f x) + when (v < w) (P.yield x) + go (min v w) + + +------------------------------- +-- Utils +------------------------------- + + +partition :: Int -> [a] -> [[a]] +partition n = + transpose . group n + where + group _ [] = [] + group k xs = take k xs : (group k $ drop k xs) diff --git a/src/Numeric/SGD/AdaDelta.hs b/src/Numeric/SGD/AdaDelta.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/SGD/AdaDelta.hs @@ -0,0 +1,97 @@ +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE DeriveGeneric #-} + + +-- | Provides the `adaDelta` function which implements the AdaDelta algorithm +-- as described in the following paper: +-- +-- * https://arxiv.org/pdf/1212.5701.pdf + + +module Numeric.SGD.AdaDelta + ( Config(..) + , adaDelta + ) where + + +import GHC.Generics (Generic) + +import Prelude hiding (div) +-- import Control.Monad (when) + +import Data.Default + +import qualified Pipes as P + +import Numeric.SGD.Type +import Numeric.SGD.ParamSet +-- import Numeric.SGD.Args + + +-- | AdaDelta configuration +data Config = Config + { decay :: Double + -- ^ Exponential decay parameter + , eps :: Double + -- ^ Epsilon value + } deriving (Show, Eq, Ord, Generic) + +instance Default Config where + def = Config + { decay = 0.9 + , eps = 1.0e-6 + } + + +-- | Perform gradient descent using the AdaDelta algorithm. +-- See "Numeric.SGD.AdaDelta" for more information. +adaDelta + :: (Monad m, ParamSet p) + => Config + -- ^ AdaDelta configuration + -> (e -> p -> p) + -- ^ Gradient on a training element + -> SGD m e p +adaDelta Config{..} gradient net0 = + + let zr = zero net0 + in go (0 :: Integer) zr zr zr net0 + + where + + go k expSqGradPrev expSqDeltaPrev deltaPrev net = do + x <- P.await + let grad = gradient x net + expSqGrad = scale decay expSqGradPrev + `add` scale (1-decay) (square grad) + rmsGrad = squareRoot (pmap (+eps) expSqGrad) + expSqDelta = scale decay expSqDeltaPrev + `add` scale (1-decay) (square deltaPrev) + rmsDelta = squareRoot (pmap (+eps) expSqDelta) + delta = (rmsDelta `mul` grad) `div` rmsGrad + newNet = net `sub` delta + P.yield newNet + go (k+1) expSqGrad expSqDelta delta newNet + + +------------------------------- +-- Utils +------------------------------- + + +-- | Scaling +scale :: ParamSet p => Double -> p -> p +scale x = pmap (*x) +{-# INLINE scale #-} + + +-- | Root square +squareRoot :: ParamSet p => p -> p +squareRoot = pmap sqrt +{-# INLINE squareRoot #-} + + +-- | Square +square :: ParamSet p => p -> p +square x = x `mul` x +{-# INLINE square #-} diff --git a/src/Numeric/SGD/Adam.hs b/src/Numeric/SGD/Adam.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/SGD/Adam.hs @@ -0,0 +1,126 @@ +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE DeriveGeneric #-} + + +-- | Provides the `adam` function which implements the Adam algorithm based on +-- the paper: +-- +-- * https://arxiv.org/pdf/1412.6980 + + +module Numeric.SGD.Adam + ( Config(..) + , scaleTau + , adam + ) where + + +import GHC.Generics (Generic) + +import Prelude hiding (div) +-- import Control.Monad (when) + +import Data.Default + +import qualified Pipes as P + +import Numeric.SGD.Type +import Numeric.SGD.ParamSet + +-- import Debug.Trace (trace) + + +-- | AdaDelta configuration +data Config = Config + { alpha0 :: Double + -- ^ Initial step size + , tau :: Double + -- ^ The step size after k * `tau` iterations = `alpha0` / (k + 1) + , beta1 :: Double + -- ^ 1st exponential moment decay + , beta2 :: Double + -- ^ 1st exponential moment decay + , eps :: Double + -- ^ Epsilon + } deriving (Show, Eq, Ord, Generic) + +instance Default Config where + def = Config + { alpha0 = 0.001 + , tau = 10000 + , beta1 = 0.9 + , beta2 = 0.999 + , eps = 1.0e-8 + } + + +-- | Scale the `tau` parameter. Useful e.g. to account for the size of the +-- training dataset. +scaleTau :: Double -> Config -> Config +scaleTau coef cfg = cfg {tau = coef * tau cfg} + + +-- | Perform gradient descent using the Adam algorithm. +-- See "Numeric.SGD.Adam" for more information. +adam + :: (Monad m, ParamSet p) + => Config + -- ^ Adam configuration + -> (e -> p -> p) + -- ^ Gradient on a training element + -> SGD m e p +adam Config{..} gradient net0 = + + let zr = zero net0 + in go (1 :: Integer) zr zr net0 + + where + + -- Gain in the k-th iteration + alpha k + = (alpha0 * tau) + / (tau + fromIntegral k) + +-- report t action = +-- if t `mod` 25 == 0 +-- then trace (show (tau, t, alpha t)) action +-- else action + + go t m v net = do + -- x <- report t (P.await) + x <- P.await + let g = gradient x net + m' = pmap (*beta1) m `add` pmap (*(1-beta1)) g + v' = pmap (*beta2) v `add` pmap (*(1-beta2)) (g `mul` g) + -- bias-corrected moment estimates + mb = pmap (/(1-beta1^t)) m' + vb = pmap (/(1-beta2^t)) v' + newNet = net `sub` + ( pmap (*alpha t) mb `div` + (pmap (+eps) (pmap sqrt vb)) + ) + newNet `seq` P.yield newNet + go (t+1) m' v' newNet + + +------------------------------- +-- Utils +------------------------------- + + +-- -- | Scaling +-- scale :: ParamSet p => Double -> p -> p +-- scale x = pmap (*x) +-- {-# INLINE scale #-} +-- +-- +-- -- | Root square +-- squareRoot :: ParamSet p => p -> p +-- squareRoot = pmap sqrt +-- {-# INLINE squareRoot #-} +-- +-- +-- -- | Square +-- square :: ParamSet p => p -> p +-- square x = x `mul` x +-- {-# INLINE square #-} diff --git a/src/Numeric/SGD/DataSet.hs b/src/Numeric/SGD/DataSet.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/SGD/DataSet.hs @@ -0,0 +1,149 @@ +{-# LANGUAGE RecordWildCards #-} + + +-- | Provides the `DataSet` type which abstracts over the actual (IO-based) +-- representation of the training dataset. + + +module Numeric.SGD.DataSet +( +-- * Dataset + DataSet (..) +, shuffle +-- * Reading +, loadData +, randomSample +-- * Construction +, withVect +, withDisk +-- , withData +) where + + +import Control.Monad (forM_) +import qualified Control.Monad.State.Strict as S + +import System.IO.Temp (withTempDirectory) +import System.IO.Unsafe (unsafeInterleaveIO) +import System.FilePath (()) +import qualified System.Random as R +import System.Random.Shuffle (shuffleM) + +import Data.Binary (Binary, encodeFile, decode) +import qualified Data.ByteString as B +import qualified Data.ByteString.Lazy as BL +import qualified Data.Vector as V +import qualified Data.Map.Strict as M + + +------------------------------- +-- Type +------------------------------- + + +-- | Dataset stored on a disk +data DataSet elem = DataSet + { size :: Int + -- ^ The size of the dataset; the individual indices are + -- [0, 1, ..., size - 1] + , elemAt :: Int -> IO elem + -- ^ Get the dataset element with the given identifier + } + + +------------------------------------------- +-- Reading +------------------------------------------- + + +-- | Lazily load the entire dataset from a disk. +loadData :: DataSet a -> IO [a] +loadData DataSet{..} = lazyMapM elemAt [0 .. size - 1] + + +-- -- | A dataset sample of the given size. +-- sample :: R.RandomGen g => g -> Int -> DataSet a -> IO ([a], g) +-- sample g 0 _ = return ([], g) +-- sample g n dataset = do +-- (xs, g') <- sample g (n-1) dataset +-- let (i, g'') = R.next g' +-- x <- dataset `elemAt` (i `mod` size dataset) +-- return (x:xs, g'') + + +-- | Shuffle the dataset. +shuffle :: DataSet a -> IO (DataSet a) +shuffle DataSet{..} = do + let ixs = [0 .. size - 1] + ixs' <- shuffleM ixs + let m = M.fromList (zip ixs ixs') + return $ DataSet + { size = size + , elemAt = elemAt . (m M.!) + } + + +-- | Random dataset sample with a specified number of elements (loaded eagerly) +randomSample :: Int -> DataSet a -> IO [a] +randomSample k dataSet + | k <= 0 = return [] + | otherwise = do + ix <- R.randomRIO (0, size dataSet - 1) + x <- elemAt dataSet ix + (x:) <$> randomSample (k-1) dataSet + + +------------------------------------------- +-- Construction +------------------------------------------- + + +-- | Construct dataset from a list of elements, store it as a vector, and run +-- the given handler. +withVect :: [a] -> (DataSet a -> IO b) -> IO b +withVect xs handler = + handler dataset + where + v = V.fromList xs + dataset = DataSet + { size = V.length v + , elemAt = \k -> return (v V.! k) } + + +-- | Construct dataset from a list of elements, store it on a disk and run the +-- given handler. Training elements must have the `Binary` instance for this +-- function to work. +withDisk :: Binary a => [a] -> (DataSet a -> IO b) -> IO b +withDisk xs handler = withTempDirectory "." ".sgd" $ \tmpDir -> do + -- We use state monad to compute the number of dataset elements. + n <- flip S.execStateT 0 $ forM_ (zip xs [0 :: Int ..]) $ \(x, ix) -> do + S.lift $ encodeFile (tmpDir show ix) x + S.modify (+1) + + -- Avoid decodeFile laziness when using some older versions of the binary + -- library (as of year 2019, this could be probably simplified) + let at ix = do + cs <- B.readFile (tmpDir show ix) + return . decode $ BL.fromChunks [cs] + + handler $ DataSet {size = n, elemAt = at} + + +------------------------------------------- +-- Lazy IO Utils +------------------------------------------- + + +-- | Lazily evaluate each action in the sequence from left to right, +-- and collect the results. +lazySequence :: [IO a] -> IO [a] +lazySequence (mx:mxs) = do + x <- mx + xs <- unsafeInterleaveIO (lazySequence mxs) + return (x : xs) +lazySequence [] = return [] + + +-- | `lazyMapM` f is equivalent to `lazySequence` . `map` f. +lazyMapM :: (a -> IO b) -> [a] -> IO [b] +lazyMapM f = lazySequence . map f diff --git a/src/Numeric/SGD/Dataset.hs b/src/Numeric/SGD/Dataset.hs deleted file mode 100644 --- a/src/Numeric/SGD/Dataset.hs +++ /dev/null @@ -1,103 +0,0 @@ -{-# LANGUAGE RecordWildCards #-} - - --- | Dataset abstraction. - - -module Numeric.SGD.Dataset -( --- * Dataset - Dataset (..) --- * Reading -, loadData -, sample --- * Construction -, withVect -, withDisk -, withData -) where - - -import Control.Monad (forM_) -import Data.Binary (Binary, encodeFile, decode) -import qualified Data.ByteString as B -import qualified Data.ByteString.Lazy as BL -import System.IO.Temp (withTempDirectory) -import System.FilePath (()) -import qualified System.Random as R -import qualified Data.Vector as V -import qualified Control.Monad.LazyIO as LazyIO -import qualified Control.Monad.State.Strict as S - - --- | A dataset with elements of type @a@. -data Dataset a = Dataset { - -- | A size of the dataset. - size :: Int - -- | Get dataset element with a given index. The set of indices - -- is of a {0, 1, .., size - 1} form. - , elemAt :: Int -> IO a } - - -------------------------------------------- --- Reading -------------------------------------------- - - --- | Lazily load dataset from a disk. -loadData :: Dataset a -> IO [a] -loadData Dataset{..} = LazyIO.mapM elemAt [0 .. size - 1] - - --- | A dataset sample of the given size. -sample :: R.RandomGen g => g -> Int -> Dataset a -> IO ([a], g) -sample g 0 _ = return ([], g) -sample g n dataset = do - (xs, g') <- sample g (n-1) dataset - let (i, g'') = R.next g' - x <- dataset `elemAt` (i `mod` size dataset) - return (x:xs, g'') - - -------------------------------------------- --- Construction -------------------------------------------- - - --- | Construct dataset from a vector of elements and run the --- given handler. -withVect :: [a] -> (Dataset a -> IO b) -> IO b -withVect xs handler = - handler dataset - where - v = V.fromList xs - dataset = Dataset - { size = V.length v - , elemAt = \k -> return (v V.! k) } - - --- | Construct dataset from a list of elements, store it on a disk --- and run the given handler. -withDisk :: Binary a => [a] -> (Dataset a -> IO b) -> IO b -withDisk xs handler = withTempDirectory "." ".sgd" $ \tmpDir -> do - -- We use state monad to compute the number of dataset elements. - n <- flip S.execStateT 0 $ forM_ (zip xs [0 :: Int ..]) $ \(x, ix) -> do - S.lift $ encodeFile (tmpDir show ix) x - S.modify (+1) - - -- We need to avoid decodeFile laziness when using some older - -- versions of the binary library. - let at ix = do - cs <- B.readFile (tmpDir show ix) - return . decode $ BL.fromChunks [cs] - - handler $ Dataset {size = n, elemAt = at} - - --- | Use disk or vector dataset representation depending on --- the first argument: when `True`, use `withDisk`, otherwise --- use `withVect`. -withData :: Binary a => Bool -> [a] -> (Dataset a -> IO b) -> IO b -withData x = case x of - True -> withDisk - False -> withVect diff --git a/src/Numeric/SGD/Grad.hs b/src/Numeric/SGD/Grad.hs deleted file mode 100644 --- a/src/Numeric/SGD/Grad.hs +++ /dev/null @@ -1,133 +0,0 @@ -{-# LANGUAGE CPP #-} - --- | A gradient is represented by an IntMap from gradient indices --- to values. Elements with no associated values in the gradient --- are assumed to have a 0 value assigned. Such elements are --- not interesting: when adding the gradient to the vector of --- parameters, only nonzero elements are taken into account. --- --- Each value associated with a gradient position is a pair of --- positive and negative components. They are stored separately --- to ensure high accuracy of computation results. --- Besides, both positive and negative components are stored --- in a logarithmic domain. - -module Numeric.SGD.Grad -( Grad -, empty -, add -, addL -, fromList -, fromLogList -, toList -, parUnions -) where - -import Data.List (foldl') -import Control.Applicative ((<$>), (<*>)) -import Control.Monad.Par (Par, runPar, get) -#if MIN_VERSION_containers(0,4,2) -import Control.Monad.Par (spawn) -#else -import Control.DeepSeq (deepseq) -import Control.Monad.Par (spawn_) -#endif -#if MIN_VERSION_containers(0,5,0) -import qualified Data.IntMap.Strict as M -#else -import qualified Data.IntMap as M -#endif - -import Numeric.SGD.LogSigned - --- | Gradient with nonzero values stored in a logarithmic domain. --- Since values equal to zero have no impact on the update phase --- of the SGD method, it is more efficient to not to store those --- components in the gradient. -type Grad = M.IntMap LogSigned - -{-# INLINE insertWith #-} -insertWith :: (a -> a -> a) -> M.Key -> a -> M.IntMap a -> M.IntMap a -#if MIN_VERSION_containers(0,5,0) -insertWith = M.insertWith -#elif MIN_VERSION_containers(0,4,1) -insertWith = M.insertWith' -#else -insertWith f k x m = - M.alter g k m - where - g my = case my of - Nothing -> Just x - Just y -> - let z = f x y - in z `seq` Just z -#endif - --- | Add normal-domain double to the gradient at the given position. -{-# INLINE add #-} -add :: Grad -> Int -> Double -> Grad -add grad i y = insertWith (+) i (logSigned y) grad - - --- | Add log-domain, singed number to the gradient at the given position. -{-# INLINE addL #-} -addL :: Grad -> Int -> LogSigned -> Grad -addL grad i y = insertWith (+) i y grad - --- | Construct gradient from a list of (index, value) pairs. --- All values from the list are added at respective gradient --- positions. -{-# INLINE fromList #-} -fromList :: [(Int, Double)] -> Grad -fromList = - let ins grad (i, y) = add grad i y - in foldl' ins empty - --- | Construct gradient from a list of (index, signed, log-domain number) --- pairs. All values from the list are added at respective gradient --- positions. -{-# INLINE fromLogList #-} -fromLogList :: [(Int, LogSigned)] -> Grad -fromLogList = - let ins grad (i, y) = addL grad i y - in foldl' ins empty - --- | Collect gradient components with values in normal domain. -{-# INLINE toList #-} -toList :: Grad -> [(Int, Double)] -toList = - let unLog (i, x) = (i, toNorm x) - in map unLog . M.assocs - --- | Empty gradient, i.e. with all elements set to 0. -{-# INLINE empty #-} -empty :: Grad -empty = M.empty - --- | Perform parallel unions operation on gradient list. --- Experimental version. -parUnions :: [Grad] -> Grad -parUnions [] = error "parUnions: empty list" -parUnions xs = runPar (parUnionsP xs) - --- | Parallel unions in the Par monad. -parUnionsP :: [Grad] -> Par Grad -parUnionsP [x] = return x -parUnionsP zs = do - let (xs, ys) = split zs -#if MIN_VERSION_containers(0,4,2) - xsP <- spawn (parUnionsP xs) - ysP <- spawn (parUnionsP ys) - M.unionWith (+) <$> get xsP <*> get ysP -#else - xsP <- spawn_ (parUnionsP xs) - ysP <- spawn_ (parUnionsP ys) - x <- M.unionWith (+) <$> get xsP <*> get ysP - M.elems x `deepseq` return x -#endif - where - split [] = ([], []) - split (x:[]) = ([x], []) - split (x:y:rest) = - let (xs, ys) = split rest - in (x:xs, y:ys) diff --git a/src/Numeric/SGD/LogSigned.hs b/src/Numeric/SGD/LogSigned.hs deleted file mode 100644 --- a/src/Numeric/SGD/LogSigned.hs +++ /dev/null @@ -1,85 +0,0 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE GeneralizedNewtypeDeriving #-} - --- | Module provides data type for signed log-domain calculations. - -module Numeric.SGD.LogSigned -( LogSigned (..) -, logSigned -, fromPos -, fromNeg -, toNorm -, toLogFloat -) where - -import qualified Data.Number.LogFloat as L -import Data.Function (on) -import Control.DeepSeq (NFData(..)) - --- | Signed real value in the logarithmic domain. -data LogSigned = LogSigned - { pos :: {-# UNPACK #-} !L.LogFloat -- ^ Positive component - , neg :: {-# UNPACK #-} !L.LogFloat -- ^ Negative component - } deriving Show - -instance Eq LogSigned where - (==) = (==) `on` toLogFloat - -instance Ord LogSigned where - compare = compare `on` toLogFloat - --- All fields are strict and unpacked. -instance NFData LogSigned where - rnf (LogSigned p q) = p `seq` q `seq` () - --- | Smart LogSigned constructor. -{-# INLINE logSigned #-} -logSigned :: Double -> LogSigned -logSigned x - | x > 0 = LogSigned (L.logFloat x) zero - | x < 0 = LogSigned zero (L.logFloat (-x)) - | otherwise = LogSigned zero zero - --- | Make LogSigned from a positive, log-domain number. -{-# INLINE fromPos #-} -fromPos :: L.LogFloat -> LogSigned -fromPos x = LogSigned x zero - --- | Make LogSigned from a negative, log-domain number. -{-# INLINE fromNeg #-} -fromNeg :: L.LogFloat -> LogSigned -fromNeg x = LogSigned zero x - --- | Shift LogSigned to a normal domain. -{-# INLINE toNorm #-} -toNorm :: LogSigned -> Double -toNorm (LogSigned x y) = L.fromLogFloat x - L.fromLogFloat y - --- | Change the 'LogSigned' to either negative 'Left' 'L.LogFloat' --- or positive 'Right' 'L.LogFloat'. -toLogFloat :: LogSigned -> Either L.LogFloat L.LogFloat -toLogFloat x = case signum x of - -1 -> Left $ neg x - pos x - 1 -> Right $ pos x - neg x - _ -> Right $ L.logFloat (0 :: Double) - -instance Num LogSigned where - LogSigned x y + LogSigned x' y' = - LogSigned (x + x') (y + y') - LogSigned x y * LogSigned x' y' = - LogSigned (x*x' + y*y') (x*y' + y*x') - LogSigned x y - LogSigned x' y' = - LogSigned (x + y') (y + x') - negate (LogSigned x y) = LogSigned y x - abs (LogSigned x y) - | x >= y = LogSigned x y - | otherwise = LogSigned y x - signum (LogSigned x y) - | x > y = 1 - | x < y = -1 - | otherwise = 0 - fromInteger = logSigned . fromInteger - -{-# INLINE zero #-} -zero :: L.LogFloat -zero = L.logFloat (0 :: Double) diff --git a/src/Numeric/SGD/Momentum.hs b/src/Numeric/SGD/Momentum.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/SGD/Momentum.hs @@ -0,0 +1,89 @@ +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE DeriveGeneric #-} + + +-- | Provides the `momentum` function which implements stochastic gradient +-- descent with momentum, following: +-- +-- * http://ruder.io/optimizing-gradient-descent/index.html#momentum + + +module Numeric.SGD.Momentum + ( Config(..) + , scaleTau + , momentum + ) where + + +import GHC.Generics (Generic) + +import Data.Default + +import qualified Pipes as P + +import Numeric.SGD.Type +import Numeric.SGD.ParamSet + + +-- | Momentum configuration +data Config = Config + { alpha0 :: Double + -- ^ Initial step size, used to scale the gradient + , tau :: Double + -- ^ The step size after k * `tau` iterations = `alpha0` / (k + 1) + , gamma :: Double + -- ^ Momentum term + } deriving (Show, Eq, Ord, Generic) + +instance Default Config where + def = Config + { alpha0 = 0.01 + , gamma = 0.9 + , tau = 1000 + } + + +-- | Scale the `tau` parameter. Useful e.g. to account for the size of the +-- training dataset. +scaleTau :: Double -> Config -> Config +scaleTau coef cfg = cfg {tau = coef * tau cfg} + + +-- | Stochastic gradient descent with momentum. See "Numeric.SGD.Momentum" for +-- more information. +momentum + :: (Monad m, ParamSet p) + => Config + -- ^ Momentum configuration + -> (e -> p -> p) + -- ^ Gradient on a training element + -> SGD m e p +momentum Config{..} gradient net0 = + + go (0 :: Integer) (zero net0) net0 + + where + + -- Gain in the k-th iteration + alpha k + = (alpha0 * tau) + / (tau + fromIntegral k) + + go k moment net = do + x <- P.await + let grad = scale (alpha k) (gradient x net) + moment' = scale gamma moment `add` grad + newNet = net `sub` moment' + P.yield newNet + go (k+1) moment' newNet + + +------------------------------- +-- Utils +------------------------------- + + +-- | Scaling +scale :: ParamSet p => Double -> p -> p +scale x = pmap (*x) +{-# INLINE scale #-} diff --git a/src/Numeric/SGD/ParamSet.hs b/src/Numeric/SGD/ParamSet.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/SGD/ParamSet.hs @@ -0,0 +1,505 @@ +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE EmptyCase #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE DefaultSignatures #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DeriveGeneric #-} + +-- {-# OPTIONS_GHC -O -ddump-rule-firings #-} + + +-- | Provides the class `ParamSet` which is used to represent the set of +-- parameters of a particular model. The goal of SGD is then to find the +-- parameter values which minimize a given objective function. + + +module Numeric.SGD.ParamSet + ( + -- * Class + ParamSet(..) + -- * Generics + , GPMap + , GAdd + , GSub + , GDiv + , GMul + , GNorm2 + ) where + + +import GHC.Generics +import GHC.TypeNats (KnownNat) + +import Prelude hiding (div) + +import qualified Data.Map.Strict as M + +import qualified Numeric.LinearAlgebra.Static as LA + + +-- | Class of types that can be treated as parameter sets. It provides basic +-- element-wise operations (addition, multiplication, mapping) which are +-- required to perform stochastic gradient descent. Many of the operations +-- (`add`, `mul`, `sub`, `div`, etc.) have the same interpretation and follow +-- the same laws (e.g. associativity) as the corresponding operations in `Num` +-- and `Fractional`. +-- +-- `zero` takes a parameter set as argument and "zero out"'s all its elements +-- (as in the backprop library). This allows instances for `Maybe`, `M.Map`, +-- etc., where the structure of the parameter set is dynamic. This leads to +-- the following property: +-- +-- @add (zero x) x = x@ +-- +-- However, `zero` does not have to obey @(add (zero x) y = y)@. +-- +-- A `ParamSet` can be also seen as a (structured) vector, hence `pmap` and +-- `norm_2`. The latter is not strictly necessary to perform SGD, but it is +-- useful to control the training process. +-- +-- `pmap` should obey the following law: +-- +-- @pmap id x = x@ +-- +-- If you leave the body of an instance declaration blank, GHC Generics will be +-- used to derive instances if the type has a single constructor and each field +-- is an instance of `ParamSet`. +class ParamSet a where + -- | Element-wise mapping + pmap :: (Double -> Double) -> a -> a + + -- | Zero-out all elements + zero :: a -> a + zero = pmap (const 0.0) + +-- -- | Element-wise negation +-- neg :: a -> a +-- neg = pmap (\x -> -x) + + -- | Element-wise addition + add :: a -> a -> a + -- | Elementi-wise substruction + sub :: a -> a -> a + + -- | Element-wise multiplication + mul :: a -> a -> a + -- | Element-wise division + div :: a -> a -> a + + -- | L2 norm + norm_2 :: a -> Double + +-- default zero :: (Generic a, GZero (Rep a)) => a -> a +-- zero = genericZero +-- {-# INLINE zero #-} + + default pmap + :: (Generic a, GPMap (Rep a)) + => (Double -> Double) -> a -> a + pmap = genericPMap + {-# INLINE pmap #-} + + + default add :: (Generic a, GAdd (Rep a)) => a -> a -> a + add = genericAdd + {-# INLINE add #-} + + default sub :: (Generic a, GSub (Rep a)) => a -> a -> a + sub = genericSub + {-# INLINE sub #-} + + default mul :: (Generic a, GMul (Rep a)) => a -> a -> a + mul = genericMul + {-# INLINE mul #-} + + default div :: (Generic a, GDiv (Rep a)) => a -> a -> a + div = genericDiv + {-# INLINE div #-} + + default norm_2 :: (Generic a, GNorm2 (Rep a)) => a -> Double + norm_2 = genericNorm2 + {-# INLINE norm_2 #-} + + +{-# RULES +"ParamSet pmap/pmap" forall f g p. pmap f (pmap g p) = pmap (f . g) p + #-} + + +-- {-# RULES +-- "ParamSet pmap/add/pmap" forall f g p h q. +-- pmap f (add (pmap g p) (pmap h q)) +-- = add (pmap (f . g) p) (pmap (f . h) q) +-- #-} + + +-- -- | 'add' using GHC Generics; works if all fields are instances of +-- -- 'ParamSet', but only for values with single constructors. +-- genericZero :: (Generic a, GZero (Rep a)) => a -> a +-- genericZero x = to $ gzero (from x) +-- {-# INLINE genericZero #-} + + +-- | 'add' using GHC Generics; works if all fields are instances of +-- 'ParamSet', but only for values with single constructors. +genericAdd :: (Generic a, GAdd (Rep a)) => a -> a -> a +genericAdd x y = to $ gadd (from x) (from y) +{-# INLINE genericAdd #-} + + +-- | 'sub' using GHC Generics; works if all fields are instances of +-- 'ParamSet', but only for values with single constructors. +genericSub :: (Generic a, GSub (Rep a)) => a -> a -> a +genericSub x y = to $ gsub (from x) (from y) +{-# INLINE genericSub #-} + + +-- | 'div' using GHC Generics; works if all fields are instances of +-- 'ParamSet', but only for values with single constructors. +genericDiv :: (Generic a, GDiv (Rep a)) => a -> a -> a +genericDiv x y = to $ gdiv (from x) (from y) +{-# INLINE genericDiv #-} + + +-- | 'mul' using GHC Generics; works if all fields are instances of +-- 'ParamSet', but only for values with single constructors. +genericMul :: (Generic a, GMul (Rep a)) => a -> a -> a +genericMul x y = to $ gmul (from x) (from y) +{-# INLINE genericMul #-} + + +-- | 'norm_2' using GHC Generics; works if all fields are instances of +-- 'ParamSet', but only for values with single constructors. +genericNorm2 :: (Generic a, GNorm2 (Rep a)) => a -> Double +genericNorm2 x = gnorm_2 (from x) +{-# INLINE genericNorm2 #-} + + +-- | 'pmap' using GHC Generics; works if all fields are instances of +-- 'ParamSet', but only for values with single constructors. +genericPMap :: (Generic a, GPMap (Rep a)) => (Double -> Double) -> a -> a +genericPMap f x = to $ gpmap f (from x) +{-# INLINE genericPMap #-} + + +-------------------------------------------------- +-- Generics +-- +-- Partially borrowed from the backprop library +-------------------------------------------------- + + +-- -- | Helper class for automatically deriving 'add' using GHC Generics. +-- class GZero f where +-- gzero :: f t -> f t +-- +-- instance ParamSet p => GZero (K1 i p) where +-- gzero (K1 x) = K1 (zero x) +-- {-# INLINE gzero #-} +-- +-- instance (GZero f, GZero g) => GZero (f :*: g) where +-- gzero (x1 :*: y1) = x2 :*: y2 +-- where +-- !x2 = gzero x1 +-- !y2 = gzero y1 +-- {-# INLINE gzero #-} +-- +-- instance GZero V1 where +-- gzero = \case {} +-- {-# INLINE gzero #-} +-- +-- instance GZero U1 where +-- gzero _ = U1 +-- {-# INLINE gzero #-} +-- +-- instance GZero f => GZero (M1 i c f) where +-- gzero (M1 x) = M1 (gzero x) +-- {-# INLINE gzero #-} +-- +-- -- instance GZero f => GZero (f :.: g) where +-- -- gzero = Comp1 gzero +-- -- {-# INLINE gzero #-} + + +-- | Helper class for automatically deriving 'add' using GHC Generics. +class GAdd f where + gadd :: f t -> f t -> f t + +instance ParamSet a => GAdd (K1 i a) where + gadd (K1 x) (K1 y) = K1 (add x y) + {-# INLINE gadd #-} + +instance (GAdd f, GAdd g) => GAdd (f :*: g) where + gadd (x1 :*: y1) (x2 :*: y2) = x3 :*: y3 + where + !x3 = gadd x1 x2 + !y3 = gadd y1 y2 + {-# INLINE gadd #-} + +instance GAdd V1 where + gadd = \case {} + {-# INLINE gadd #-} + +instance GAdd U1 where + gadd _ _ = U1 + {-# INLINE gadd #-} + +instance GAdd f => GAdd (M1 i c f) where + gadd (M1 x) (M1 y) = M1 (gadd x y) + {-# INLINE gadd #-} + +-- instance GAdd f => GAdd (f :.: g) where +-- gadd (Comp1 x) (Comp1 y) = Comp1 (gadd x y) +-- {-# INLINE gadd #-} + + +-- | Helper class for automatically deriving 'sub' using GHC Generics. +class GSub f where + gsub :: f t -> f t -> f t + +instance ParamSet a => GSub (K1 i a) where + gsub (K1 x) (K1 y) = K1 (sub x y) + {-# INLINE gsub #-} + +instance (GSub f, GSub g) => GSub (f :*: g) where + gsub (x1 :*: y1) (x2 :*: y2) = x3 :*: y3 + where + !x3 = gsub x1 x2 + !y3 = gsub y1 y2 + {-# INLINE gsub #-} + +instance GSub V1 where + gsub = \case {} + {-# INLINE gsub #-} + +instance GSub U1 where + gsub _ _ = U1 + {-# INLINE gsub #-} + +instance GSub f => GSub (M1 i c f) where + gsub (M1 x) (M1 y) = M1 (gsub x y) + {-# INLINE gsub #-} + +-- instance GSub f => GSub (f :.: g) where +-- gsub (Comp1 x) (Comp1 y) = Comp1 (gsub x y) +-- {-# INLINE gsub #-} + + +-- | Helper class for automatically deriving 'mul' using GHC Generics. +class GMul f where + gmul :: f t -> f t -> f t + +instance ParamSet a => GMul (K1 i a) where + gmul (K1 x) (K1 y) = K1 (mul x y) + {-# INLINE gmul #-} + +instance (GMul f, GMul g) => GMul (f :*: g) where + gmul (x1 :*: y1) (x2 :*: y2) = x3 :*: y3 + where + !x3 = gmul x1 x2 + !y3 = gmul y1 y2 + {-# INLINE gmul #-} + +instance GMul V1 where + gmul = \case {} + {-# INLINE gmul #-} + +instance GMul U1 where + gmul _ _ = U1 + {-# INLINE gmul #-} + +instance GMul f => GMul (M1 i c f) where + gmul (M1 x) (M1 y) = M1 (gmul x y) + {-# INLINE gmul #-} + +-- instance GMul f => GMul (f :.: g) where +-- gmul (Comp1 x) (Comp1 y) = Comp1 (gmul x y) +-- {-# INLINE gmul #-} + + +-- | Helper class for automatically deriving 'div' using GHC Generics. +class GDiv f where + gdiv :: f t -> f t -> f t + +instance ParamSet a => GDiv (K1 i a) where + gdiv (K1 x) (K1 y) = K1 (div x y) + {-# INLINE gdiv #-} + +instance (GDiv f, GDiv g) => GDiv (f :*: g) where + gdiv (x1 :*: y1) (x2 :*: y2) = x3 :*: y3 + where + !x3 = gdiv x1 x2 + !y3 = gdiv y1 y2 + {-# INLINE gdiv #-} + +instance GDiv V1 where + gdiv = \case {} + {-# INLINE gdiv #-} + +instance GDiv U1 where + gdiv _ _ = U1 + {-# INLINE gdiv #-} + +instance GDiv f => GDiv (M1 i c f) where + gdiv (M1 x) (M1 y) = M1 (gdiv x y) + {-# INLINE gdiv #-} + +-- instance GDiv f => GDiv (f :.: g) where +-- gdiv (Comp1 x) (Comp1 y) = Comp1 (gdiv x y) +-- {-# INLINE gdiv #-} + + +-- | Helper class for automatically deriving 'norm_2' using GHC Generics. +class GNorm2 f where + gnorm_2 :: f t -> Double + +instance ParamSet a => GNorm2 (K1 i a) where + gnorm_2 (K1 x) = norm_2 x + {-# INLINE gnorm_2 #-} + +instance (GNorm2 f, GNorm2 g) => GNorm2 (f :*: g) where + gnorm_2 (x1 :*: y1) = + sqrt ((x2 ^ (2 :: Int)) + (y2 ^ (2 :: Int))) + where + !x2 = gnorm_2 x1 + !y2 = gnorm_2 y1 + {-# INLINE gnorm_2 #-} + +instance GNorm2 V1 where + gnorm_2 = \case {} + {-# INLINE gnorm_2 #-} + +instance GNorm2 U1 where + gnorm_2 _ = 0 + {-# INLINE gnorm_2 #-} + +instance GNorm2 f => GNorm2 (M1 i c f) where + gnorm_2 (M1 x) = gnorm_2 x + {-# INLINE gnorm_2 #-} + +-- -- TODO: Make sure this makes sense +-- instance GNorm2 f => GNorm2 (f :.: g) where +-- gnorm_2 (Comp1 x) = gnorm_2 x +-- {-# INLINE gnorm_2 #-} + + +-- | Helper class for automatically deriving 'pmap' using GHC Generics. +class GPMap f where + gpmap :: (Double -> Double) -> f t -> f t + +instance ParamSet a => GPMap (K1 i a) where + gpmap f (K1 x) = K1 (pmap f x) + {-# INLINE gpmap #-} + +instance (GPMap f, GPMap g) => GPMap (f :*: g) where + gpmap f (x1 :*: y1) = x2 :*: y2 + where + !x2 = gpmap f x1 + !y2 = gpmap f y1 + {-# INLINE gpmap #-} + +instance GPMap V1 where + gpmap _ = \case {} + {-# INLINE gpmap #-} + +instance GPMap U1 where + gpmap _ _ = U1 + {-# INLINE gpmap #-} + +instance GPMap f => GPMap (M1 i c f) where + gpmap f (M1 x) = M1 (gpmap f x) + {-# INLINE gpmap #-} + +-- instance GPMap f => GPMap (f :.: g) where +-- gpmap f (Comp1 x) = Comp1 (gpmap f x) +-- {-# INLINE gpmap #-} + + +-------------------------------------------------- +-- Basic instances +-------------------------------------------------- + + +instance ParamSet Double where + zero = const 0 + pmap = id + add = (+) + sub = (-) + mul = (*) + div = (/) + norm_2 = abs + + +instance (ParamSet a, ParamSet b) => ParamSet (a, b) where + pmap f (x, y) = (pmap f x, pmap f y) + add (x1, y1) (x2, y2) = (x1 `add` x2, y1 `add` y2) + sub (x1, y1) (x2, y2) = (x1 `sub` x2, y1 `sub` y2) + mul (x1, y1) (x2, y2) = (x1 `mul` x2, y1 `mul` y2) + div (x1, y1) (x2, y2) = (x1 `div` x2, y1 `div` y2) + norm_2 (x, y) + = sqrt . sum . map ((^(2::Int))) + $ [norm_2 x, norm_2 y] + + +instance (KnownNat n) => ParamSet (LA.R n) where + zero = const 0 + pmap = LA.dvmap + add = (+) + sub = (-) + mul = (*) + div = (/) + norm_2 = LA.norm_2 + + +instance (KnownNat n, KnownNat m) => ParamSet (LA.L n m) where + zero = const 0 + pmap = LA.dmmap + add = (+) + sub = (-) + mul = (*) + div = (/) + norm_2 = LA.norm_2 + + +-- | `Nothing` represents a deactivated parameter set component. If `Nothing` +-- is given as an argument to one of the `ParamSet` operations, the result is +-- `Nothing` as well. +-- +-- This differs from the corresponding instance in the backprop library, where +-- `Nothing` is equivalent to `Just 0`. However, the implementation below +-- seems to correspond adequately enough to the notion that a particular +-- component is either active or not in both the parameter set and the +-- gradient, hence it doesn't make sense to combine `Just` with `Nothing`. +instance (ParamSet a) => ParamSet (Maybe a) where + zero = fmap zero + pmap = fmap . pmap + + add (Just x) (Just y) = Just (add x y) + add _ _ = Nothing + + sub (Just x) (Just y) = Just (sub x y) + sub _ _ = Nothing + + mul (Just x) (Just y) = Just (mul x y) + mul _ _ = Nothing + + div (Just x) (Just y) = Just (div x y) + div _ _ = Nothing + + norm_2 = maybe 0 norm_2 + + +-- | A map with different parameter sets (of the same type) assigned to the +-- individual keys. +-- +-- When combining two maps with different sets of keys, only their intersection +-- is preserved. +instance (Ord k, ParamSet a) => ParamSet (M.Map k a) where + zero = fmap zero + pmap f = fmap (pmap f) + add = M.intersectionWith add + sub = M.intersectionWith sub + mul= M.intersectionWith mul + div= M.intersectionWith div + norm_2 = sqrt . sum . map ((^(2::Int)) . norm_2) . M.elems diff --git a/src/Numeric/SGD/Sparse.hs b/src/Numeric/SGD/Sparse.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/SGD/Sparse.hs @@ -0,0 +1,167 @@ +{-# LANGUAGE RecordWildCards #-} + + +-- | Stochastic gradient descent using mutable vectors for efficient parameter +-- update. This module is intended for use with sparse features. If you use +-- dense feature vectors (as arise e.g. in deep learning), have a look at +-- "Numeric.SGD". +-- +-- Currently only the Gaussian regularization is implemented. +-- +-- SGD with momentum is known to converge faster than vanilla SGD. It's +-- implementation can be found in "Numeric.SGD.Sparse.Momentum". + + +module Numeric.SGD.Sparse +( SgdArgs (..) +, sgdArgsDefault +, Para +, sgd +, module Numeric.SGD.Sparse.Grad +, module Numeric.SGD.DataSet +) where + + +import Control.Monad (forM_) +-- import qualified System.Random as R +import qualified Data.Vector.Unboxed as U +import qualified Data.Vector.Unboxed.Mutable as UM +import qualified Control.Monad.Primitive as Prim + +import Numeric.SGD.Sparse.Grad +import Numeric.SGD.DataSet + + +-- | SGD parameters controlling the learning process. +data SgdArgs = SgdArgs + { -- | Size of the batch + batchSize :: Int + -- | Regularization variance + , regVar :: Double + -- | Number of iterations + , iterNum :: Double + -- | Initial gain parameter + , gain0 :: Double + -- | After how many iterations over the entire dataset + -- the gain parameter is halved + , tau :: Double } + + +-- | Default SGD parameter values. +sgdArgsDefault :: SgdArgs +sgdArgsDefault = SgdArgs + { batchSize = 30 + , regVar = 10 + , iterNum = 10 + , gain0 = 1 + , tau = 5 } + + +-- | Vector of parameters. +type Para = U.Vector Double + + +-- | Type synonym for mutable vector with Double values. +type MVect = UM.MVector (Prim.PrimState IO) Double + + +-- | A stochastic gradient descent method. +-- A notification function can be used to provide user with +-- information about the progress of the learning. +sgd + :: SgdArgs -- ^ SGD parameter values + -> (Para -> Int -> IO ()) -- ^ Notification run every update + -> (Para -> x -> Grad) -- ^ Gradient for dataset element + -> DataSet x -- ^ Dataset + -> Para -- ^ Starting point + -> IO Para -- ^ SGD result +sgd SgdArgs{..} notify mkGrad dataset x0 = do + u <- UM.new (U.length x0) + -- doIt u 0 (R.mkStdGen 0) =<< U.thaw x0 + doIt u 0 =<< U.thaw x0 + where + -- Gain in k-th iteration. + gain k = (gain0 * tau) / (tau + done k) + + -- Number of completed iterations over the full dataset. + done :: Int -> Double + done k + = fromIntegral (k * batchSize) + / fromIntegral (size dataset) + -- doneTotal :: Int -> Int + -- doneTotal = floor . done + + -- Regularization (Guassian prior) + regularization k = regCoef + where + regCoef = (1.0 - gain k * iVar) ** coef + iVar = 1.0 / regVar + coef = fromIntegral batchSize + / fromIntegral (size dataset) + +-- -- Regularization (Guassian prior) after a full dataset pass +-- regularization k = 1.0 - (gain k / regVar) + + doIt u k x + | done k > iterNum = do + frozen <- U.unsafeFreeze x + notify frozen k + return frozen + | otherwise = do + -- (batch, stdGen') <- sample stdGen batchSize dataset + batch <- randomSample batchSize dataset + + -- Regularization + -- when (doneTotal (k - 1) /= doneTotal k) $ do + -- <- we now apply regularization each step rather than each + -- dataset pass + let regParam = regularization k + -- putStrLn $ "\nApplying regularization (params *= " ++ show regParam ++ ")" + scale regParam x + +-- -- Regularization +-- when (doneTotal (k - 1) /= doneTotal k) $ do +-- let regParam = regularization k +-- putStrLn $ "\nApplying regularization (params *= " ++ show regParam ++ ")" +-- scale regParam x + + -- Freeze mutable vector of parameters. The frozen version is + -- then supplied to external mkGrad function provided by user. + frozen <- U.unsafeFreeze x + notify frozen k + + -- let grad = M.unions (map (mkGrad frozen) batch) + let grad = parUnions (map (mkGrad frozen) batch) + addUp grad u + scale (gain k) u + + x' <- U.unsafeThaw frozen + u `addTo` x' + doIt u (k+1) x' + + +-- | Add up all gradients and store results in normal domain. +addUp :: Grad -> MVect -> IO () +addUp grad v = do + UM.set v 0 + forM_ (toList grad) $ \(i, x) -> do + y <- UM.unsafeRead v i + UM.unsafeWrite v i (x + y) + + +-- | Scale the vector by the given value. +scale :: Double -> MVect -> IO () +scale c v = do + forM_ [0 .. UM.length v - 1] $ \i -> do + y <- UM.unsafeRead v i + UM.unsafeWrite v i (c * y) + + +-- | Apply gradient to the parameters vector, that is add the first vector to +-- the second one. +addTo :: MVect -> MVect -> IO () +addTo w v = do + forM_ [0 .. UM.length v - 1] $ \i -> do + x <- UM.unsafeRead v i + y <- UM.unsafeRead w i + UM.unsafeWrite v i (x + y) diff --git a/src/Numeric/SGD/Sparse/Grad.hs b/src/Numeric/SGD/Sparse/Grad.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/SGD/Sparse/Grad.hs @@ -0,0 +1,131 @@ +{-# LANGUAGE CPP #-} + +-- | A gradient is represented by an IntMap from gradient indices to values. +-- Elements with no associated values in the gradient are assumed to have a 0 +-- value assigned. Such elements are of no interest: when adding the gradient +-- to the vector of parameters, only non-zero elements are taken into account. +-- +-- Each value associated with a gradient position is a pair of positive and +-- negative components. They are stored separately to ensure high accuracy of +-- computation results. Besides, both positive and negative components are +-- stored in a logarithmic domain. + +module Numeric.SGD.Sparse.Grad +( Grad +, empty +, add +, addL +, fromList +, fromLogList +, toList +, parUnions +) where + +import Data.List (foldl') +import Control.Applicative ((<$>), (<*>)) +import Control.Monad.Par (Par, runPar, get) +#if MIN_VERSION_containers(0,4,2) +import Control.Monad.Par (spawn) +#else +import Control.DeepSeq (deepseq) +import Control.Monad.Par (spawn_) +#endif +#if MIN_VERSION_containers(0,5,0) +import qualified Data.IntMap.Strict as M +#else +import qualified Data.IntMap as M +#endif + +import Numeric.SGD.Sparse.LogSigned + +-- | Gradient with nonzero values stored in a logarithmic domain. +-- Since values equal to zero have no impact on the update phase +-- of the SGD method, it is more efficient to not to store those +-- components in the gradient. +type Grad = M.IntMap LogSigned + +{-# INLINE insertWith #-} +insertWith :: (a -> a -> a) -> M.Key -> a -> M.IntMap a -> M.IntMap a +#if MIN_VERSION_containers(0,5,0) +insertWith = M.insertWith +#elif MIN_VERSION_containers(0,4,1) +insertWith = M.insertWith' +#else +insertWith f k x m = + M.alter g k m + where + g my = case my of + Nothing -> Just x + Just y -> + let z = f x y + in z `seq` Just z +#endif + +-- | Add normal-domain double to the gradient at the given position. +{-# INLINE add #-} +add :: Grad -> Int -> Double -> Grad +add grad i y = insertWith (+) i (logSigned y) grad + + +-- | Add log-domain, singed number to the gradient at the given position. +{-# INLINE addL #-} +addL :: Grad -> Int -> LogSigned -> Grad +addL grad i y = insertWith (+) i y grad + +-- | Construct gradient from a list of (index, value) pairs. +-- All values from the list are added at respective gradient +-- positions. +{-# INLINE fromList #-} +fromList :: [(Int, Double)] -> Grad +fromList = + let ins grad (i, y) = add grad i y + in foldl' ins empty + +-- | Construct gradient from a list of (index, signed, log-domain number) +-- pairs. All values from the list are added at respective gradient +-- positions. +{-# INLINE fromLogList #-} +fromLogList :: [(Int, LogSigned)] -> Grad +fromLogList = + let ins grad (i, y) = addL grad i y + in foldl' ins empty + +-- | Collect gradient components with values in normal domain. +{-# INLINE toList #-} +toList :: Grad -> [(Int, Double)] +toList = + let unLog (i, x) = (i, toNorm x) + in map unLog . M.assocs + +-- | Empty gradient, i.e. with all elements set to 0. +{-# INLINE empty #-} +empty :: Grad +empty = M.empty + +-- | Perform parallel unions operation on gradient list. +-- Experimental version. +parUnions :: [Grad] -> Grad +parUnions [] = error "parUnions: empty list" +parUnions xs = runPar (parUnionsP xs) + +-- | Parallel unions in the Par monad. +parUnionsP :: [Grad] -> Par Grad +parUnionsP [x] = return x +parUnionsP zs = do + let (xs, ys) = split zs +#if MIN_VERSION_containers(0,4,2) + xsP <- spawn (parUnionsP xs) + ysP <- spawn (parUnionsP ys) + M.unionWith (+) <$> get xsP <*> get ysP +#else + xsP <- spawn_ (parUnionsP xs) + ysP <- spawn_ (parUnionsP ys) + x <- M.unionWith (+) <$> get xsP <*> get ysP + M.elems x `deepseq` return x +#endif + where + split [] = ([], []) + split (x:[]) = ([x], []) + split (x:y:rest) = + let (xs, ys) = split rest + in (x:xs, y:ys) diff --git a/src/Numeric/SGD/Sparse/LogSigned.hs b/src/Numeric/SGD/Sparse/LogSigned.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/SGD/Sparse/LogSigned.hs @@ -0,0 +1,85 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} + +-- | Module provides data type for signed log-domain calculations. + +module Numeric.SGD.Sparse.LogSigned +( LogSigned (..) +, logSigned +, fromPos +, fromNeg +, toNorm +, toLogFloat +) where + +import qualified Data.Number.LogFloat as L +import Data.Function (on) +import Control.DeepSeq (NFData(..)) + +-- | Signed real value in the logarithmic domain. +data LogSigned = LogSigned + { pos :: {-# UNPACK #-} !L.LogFloat -- ^ Positive component + , neg :: {-# UNPACK #-} !L.LogFloat -- ^ Negative component + } deriving Show + +instance Eq LogSigned where + (==) = (==) `on` toLogFloat + +instance Ord LogSigned where + compare = compare `on` toLogFloat + +-- All fields are strict and unpacked. +instance NFData LogSigned where + rnf (LogSigned p q) = p `seq` q `seq` () + +-- | Smart LogSigned constructor. +{-# INLINE logSigned #-} +logSigned :: Double -> LogSigned +logSigned x + | x > 0 = LogSigned (L.logFloat x) zero + | x < 0 = LogSigned zero (L.logFloat (-x)) + | otherwise = LogSigned zero zero + +-- | Make LogSigned from a positive, log-domain number. +{-# INLINE fromPos #-} +fromPos :: L.LogFloat -> LogSigned +fromPos x = LogSigned x zero + +-- | Make LogSigned from a negative, log-domain number. +{-# INLINE fromNeg #-} +fromNeg :: L.LogFloat -> LogSigned +fromNeg x = LogSigned zero x + +-- | Shift LogSigned to a normal domain. +{-# INLINE toNorm #-} +toNorm :: LogSigned -> Double +toNorm (LogSigned x y) = L.fromLogFloat x - L.fromLogFloat y + +-- | Change the 'LogSigned' to either negative 'Left' 'L.LogFloat' +-- or positive 'Right' 'L.LogFloat'. +toLogFloat :: LogSigned -> Either L.LogFloat L.LogFloat +toLogFloat x = case signum x of + -1 -> Left $ neg x - pos x + 1 -> Right $ pos x - neg x + _ -> Right $ L.logFloat (0 :: Double) + +instance Num LogSigned where + LogSigned x y + LogSigned x' y' = + LogSigned (x + x') (y + y') + LogSigned x y * LogSigned x' y' = + LogSigned (x*x' + y*y') (x*y' + y*x') + LogSigned x y - LogSigned x' y' = + LogSigned (x + y') (y + x') + negate (LogSigned x y) = LogSigned y x + abs (LogSigned x y) + | x >= y = LogSigned x y + | otherwise = LogSigned y x + signum (LogSigned x y) + | x > y = 1 + | x < y = -1 + | otherwise = 0 + fromInteger = logSigned . fromInteger + +{-# INLINE zero #-} +zero :: L.LogFloat +zero = L.logFloat (0 :: Double) diff --git a/src/Numeric/SGD/Sparse/Momentum.hs b/src/Numeric/SGD/Sparse/Momentum.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/SGD/Sparse/Momentum.hs @@ -0,0 +1,201 @@ +{-# LANGUAGE RecordWildCards #-} + + +-- | A version of `Numeric.SGD.Sparse` extended with momentum. + + +module Numeric.SGD.Sparse.Momentum +( SgdArgs (..) +, sgdArgsDefault +, Para +, sgd +, module Numeric.SGD.Sparse.Grad +, module Numeric.SGD.DataSet +) where + + +import Control.Monad (forM_) +-- import qualified System.Random as R +import qualified Data.Vector.Unboxed as U +import qualified Data.Vector.Unboxed.Mutable as UM +import qualified Control.Monad.Primitive as Prim + +import Numeric.SGD.Sparse.Grad +import Numeric.SGD.DataSet + + +-- | SGD parameters controlling the learning process. +data SgdArgs = SgdArgs + { -- | Size of the batch + batchSize :: Int + -- | Regularization variance + , regVar :: Double + -- | Number of iterations + , iterNum :: Double + -- | Initial gain parameter + , gain0 :: Double + -- | After how many iterations over the entire dataset + -- the gain parameter is halved + , tau :: Double + } + + +-- | Default SGD parameter values. +sgdArgsDefault :: SgdArgs +sgdArgsDefault = SgdArgs + { batchSize = 50 + , regVar = 10 + , iterNum = 10 + , gain0 = 0.25 + -- Without momentum I would rather go for '1', but with momentum the + -- gradient gets significantly larger. + , tau = 5 + } + + +-- | The gamma parameter which drives momentum. +-- +-- TODO: put in SgdArgs. +-- +gamma :: Double +gamma = 0.9 + + +-- | Vector of parameters. +type Para = U.Vector Double + + +-- | Type synonym for mutable vector with Double values. +type MVect = UM.MVector (Prim.PrimState IO) Double + + +-- | A stochastic gradient descent method. +-- A notification function can be used to provide user with +-- information about the progress of the learning. +sgd + :: SgdArgs -- ^ SGD parameter values + -> (Para -> Int -> IO ()) -- ^ Notification run every update + -> (Para -> x -> Grad) -- ^ Gradient for dataset element + -> DataSet x -- ^ DataSet + -> Para -- ^ Starting point + -> IO Para -- ^ SGD result +sgd SgdArgs{..} notify mkGrad dataset x0 = do + + putStrLn $ "Running momentum!" + + -- A vector for the momentum gradient + momentum <- UM.new (U.length x0) + + -- A worker vector for computing the actual gradients + u <- UM.new (U.length x0) + + doIt momentum u 0 =<< U.thaw x0 + + where + -- Gain in k-th iteration. + gain k = (gain0 * tau) / (tau + done k) + + -- Number of completed iterations over the full dataset. + done :: Int -> Double + done k + = fromIntegral (k * batchSize) + / fromIntegral (size dataset) + + -- Regularization (Guassian prior) parameter + regularizationParam = regCoef + where + regCoef = iVar ** coef + iVar = 1.0 / regVar + coef = fromIntegral (size dataset) + / fromIntegral batchSize + + doIt momentum u k x + + | done k > iterNum = do + frozen <- U.unsafeFreeze x + notify frozen k + return frozen + + | otherwise = do + + -- Sample the dataset + batch <- randomSample batchSize dataset + + -- NEW: comment out + -- -- Apply regularization to the parameters vector. + -- scale (regularization k) x + + -- Freeze mutable vector of parameters. The frozen version is + -- then supplied to external mkGrad function provided by user. + frozen <- U.unsafeFreeze x + notify frozen k + + -- Compute the gradient and put it in `u` + let grad = parUnions (map (mkGrad frozen) batch) + addUp grad u + + -- Apply regularization to `u` + applyRegularization regularizationParam x u + + -- Scale the gradient + scale (gain k) u + + -- Compute the new momentum + updateMomentum gamma momentum u + + x' <- U.unsafeThaw frozen + momentum `addTo` x' + doIt momentum u (k+1) x' + + +-- | Compute the new momentum (gradient) vector. +applyRegularization + :: Double -- ^ Regularization parameter + -> MVect -- ^ The parameters + -> MVect -- ^ The current gradient + -> IO () +applyRegularization regParam params grad = do + forM_ [0 .. UM.length grad - 1] $ \i -> do + x <- UM.unsafeRead grad i + y <- UM.unsafeRead params i + UM.unsafeWrite grad i $ x - regParam * y + + +-- | Compute the new momentum (gradient) vector. +updateMomentum + :: Double -- ^ The gamma parameter + -> MVect -- ^ The previous momentum + -> MVect -- ^ The scaled current gradient + -> IO () +updateMomentum gammaCoef momentum grad = do + forM_ [0 .. UM.length momentum - 1] $ \i -> do + x <- UM.unsafeRead momentum i + y <- UM.unsafeRead grad i + UM.unsafeWrite momentum i (gammaCoef * x + y) + + +-- | Add up all gradients and store results in normal domain. +addUp :: Grad -> MVect -> IO () +addUp grad v = do + UM.set v 0 + forM_ (toList grad) $ \(i, x) -> do + y <- UM.unsafeRead v i + UM.unsafeWrite v i (x + y) + + +-- | Scale the vector by the given value. +scale :: Double -> MVect -> IO () +scale c v = do + forM_ [0 .. UM.length v - 1] $ \i -> do + y <- UM.unsafeRead v i + UM.unsafeWrite v i (c * y) + + +-- | Apply gradient to the parameters vector, that is add the first vector to +-- the second one. +addTo :: MVect -> MVect -> IO () +addTo w v = do + forM_ [0 .. UM.length v - 1] $ \i -> do + x <- UM.unsafeRead v i + y <- UM.unsafeRead w i + UM.unsafeWrite v i (x + y) diff --git a/src/Numeric/SGD/Type.hs b/src/Numeric/SGD/Type.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/SGD/Type.hs @@ -0,0 +1,15 @@ +-- | Provides the basic `SGD` pipe type. + + +module Numeric.SGD.Type + ( SGD + ) where + + +import Pipes as P + + +-- | SGD is a pipe which, given the initial parameter values, consumes training +-- elements of type @e@ and outputs the subsequently calculated parameter sets +-- of type @p@. +type SGD m e p = p -> P.Pipe e p m () diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,85 @@ +{-# LANGUAGE NoMonomorphismRestriction #-} + + +import Data.Ord (Ord(..)) + +import qualified Numeric.AD as AD + +import Test.Tasty (TestTree, testGroup) +import qualified Test.Tasty as Tasty +import qualified Test.Tasty.HUnit as U +import Test.Tasty.HUnit ((@?=)) +-- import qualified Test.Tasty.SmallCheck as SC + +import qualified Numeric.SGD as SGD +import qualified Numeric.SGD.Adam as Adam +-- import qualified Numeric.SGD.AdaDelta as Ada + + +main :: IO () +main = Tasty.defaultMain tests + + +tests :: TestTree +tests = testGroup "Tests" [unitTests] + + +unitTests = testGroup "Unit tests" + [ U.testCase "Simple optimization" $ do + U.assertBool "Momentum" . approxWith 0.01 4.18 $ + sgdWith (SGD.momentum SGD.def id) + U.assertBool "Adam" . approxWith 0.01 4.18 $ + let cfg = SGD.def {Adam.alpha0 = 0.1} + in sgdWith (SGD.adam cfg id) + U.assertBool "AdaDelta" . approxWith 0.25 4.18 $ + sgdWith (SGD.adaDelta SGD.def id) +-- -- the following test does not hold +-- , U.testCase "List comparison (same length)" $ +-- [1, 2, 3] `compare` [1,2,2] @?= LT + ] + + +-------------------------------------------------- +-- Main unit test +-------------------------------------------------- + + +-- | The component objective functions +funs :: [Double -> Double] +funs = + [ \x -> 0.3*x^2 + , \x -> -2*x + , const 3 + , sin + ] + + +-- | The corresponding derivatives +derivs :: [Double -> Double] +derivs = + [ AD.diff $ \x -> 0.3*x^2 + , AD.diff $ \x -> -2*x + , AD.diff $ const 3 + , AD.diff $ sin + ] + + +-- | The total objective is the sum of the objective component functions +objective :: Double -> Double +objective x = sum $ map ($x) funs + + +-- | Perform SGD with the given SGD variant. +sgdWith typ = SGD.run typ (take 10000 $ cycle derivs) (0.0 :: Double) + + +-------------------------------------------------- +-- Utils +-------------------------------------------------- + + +-- | Is the second argument approximately equaly to the third one? +-- The first argument is the epsilon. +approxWith :: (Ord a, Floating a) => a -> a -> a -> Bool +approxWith eps x y = + x >= y - eps && x <= y + eps