module MachineLearning.NeuralNetwork.Topology
(
Topology
, LossFunc(..)
, makeTopology
, loss
, propagateForward
, propagateBackward
, numberOutputs
, initializeTheta
, initializeThetaIO
, initializeThetaM
, flatten
, unflatten
)
where
import Control.Monad (zipWithM)
import Data.List (foldl')
import qualified Control.Monad.Random as RndM
import qualified Data.Vector.Storable as V
import qualified Numeric.LinearAlgebra as LA
import MachineLearning.Types (R, Vector, Matrix)
import MachineLearning.Utils (listOfTuplesToList)
import MachineLearning.NeuralNetwork.Layer (Layer(..), Cache(..))
import MachineLearning.NeuralNetwork.Regularization (Regularization, forwardReg, backwardReg)
type LossFunc = Matrix -> Matrix -> R
data Topology = Topology [(Int, Int)] [Layer] LossFunc
makeTopology :: Int -> [Layer] -> Layer -> LossFunc -> Topology
makeTopology nInputs hiddenLayers outputLayer lossFunc =
let layers = hiddenLayers ++ [outputLayer]
layerSizes = nInputs : (map lUnits layers)
sizes = getThetaSizes layerSizes
in Topology sizes layers lossFunc
loss :: Topology -> Regularization -> Matrix -> [(Matrix, Matrix)] -> Matrix -> R
loss (Topology _ _ lf) reg x weights y =
let lossValue = lf x y
regValue = forwardReg reg weights
in lossValue + regValue
propagateForward :: Topology -> Matrix -> [(Matrix, Matrix)] -> (Matrix, [Cache])
propagateForward (Topology _ layers _) x thetaList = foldl' f (x, []) $ zip thetaList layers
where f (a, cs) (theta, hl) =
let (a', cache) = forwardPass hl a theta
in (a', cache:cs)
forwardPass :: Layer -> Matrix -> (Matrix, Matrix) -> (Matrix, Cache)
forwardPass layer a (b, w) = (a', Cache z a w)
where z = lForward layer a b w
a' = lActivation layer z
propagateBackward :: Topology -> Regularization -> Matrix -> [Cache] -> Matrix -> [(Matrix, Matrix)]
propagateBackward (Topology _ layers _) reg scores (cache:cacheList) y = gradientList
where cache' = Cache scores (cacheX cache) (cacheW cache)
cacheList' = cache':cacheList
gradientList = snd $ foldl' f (y, []) $ zip cacheList' $ reverse layers
f (da, grads) (cache, hl) =
let (da', db, dw) = backwardPass hl reg da cache
in (da', (db, dw):grads)
backwardPass :: Layer -> Regularization -> Matrix -> Cache -> (Matrix, Matrix, Matrix)
backwardPass layer reg da cache = (da', db, dw')
where delta = lActivationGradient layer (cacheZ cache) da
(da', db, dw) = lBackward layer delta cache
dw' = dw + (backwardReg reg (cacheW cache))
numberOutputs :: Topology -> Int
numberOutputs (Topology nnt _ _) = fst $ last nnt
getThetaSizes :: [Int] -> [(Int, Int)]
getThetaSizes nn = zipWith (\r c -> (r, c)) (tail nn) nn
initializeTheta :: Int -> Topology -> Vector
initializeTheta seed topology = RndM.evalRand (initializeThetaM topology) gen
where gen = RndM.mkStdGen seed
initializeThetaIO :: Topology -> IO Vector
initializeThetaIO = RndM.evalRandIO . initializeThetaM
initializeThetaM :: RndM.RandomGen g => Topology -> RndM.Rand g Vector
initializeThetaM topology = flatten <$> initializeThetaListM topology
initializeThetaListM :: RndM.RandomGen g => Topology -> RndM.Rand g [(Matrix, Matrix)]
initializeThetaListM (Topology sizes layers _) = zipWithM lInitializeThetaM layers sizes
flatten :: [(Matrix, Matrix)] -> Vector
flatten ms = V.concat $ map LA.flatten $ listOfTuplesToList ms
unflatten :: Topology -> Vector -> [(Matrix, Matrix)]
unflatten (Topology sizes _ _) v =
let offsets = reverse $ foldl' (\os (r, c) -> (r+r*c + head os):os) [0] (init sizes)
ms = zipWith (\o (r, c) -> (LA.reshape r (slice o r), LA.reshape c (slice (o+r) (r*c)))) offsets sizes
slice o n = V.slice o n v
in ms