differentiable-programming.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeApplications #-}


import Data.NumInstances.Tuple
import GHC.TypeLits
import Numeric.Backprop
import Numeric.LinearAlgebra.Static (randn, randomVector, RandDist(..), vec2)
import Numeric.LinearAlgebra.Static.Backprop hiding (vec2)


-------------------------------------
-- 3.1 自動微分
-------------------------------------


data D a = D a a


dual :: D a -> a
dual (D _ a) = a


lift :: Num a => a -> D a
lift a = D a 0


partial :: Num a => a -> D a
partial a = D a 1


instance Num a => Num (D a) where
  (D x x') + (D y y') = D (x + y) (x' + y')
  (D x x') * (D y y') = D (x * y) (x' * y + x * y')
  negate (D x x')     = D (negate x) (negate x')
  abs (D x x')        = D (abs x) (x' * (signum x))
  signum (D x x')     = D (signum x) 0
  fromInteger n       = D (fromInteger n) 0


instance Fractional a => Fractional (D a) where
  recip (D x x') = D (recip x) (-1 * x' * (recip (x * x)))
  fromRational x = D (fromRational x) 0


instance Floating a => Floating (D a) where
  pi             = D pi 0
  exp (D x x')   = D (exp x) (x' * exp x)
  log (D x x')   = D (log x) (x' / x)
  sin (D x x')   = D (sin x) (x' * cos x)
  cos (D x x')   = D (cos x) (- x' * sin x)
  asin (D x x')  = D (asin x) (x' / (sqrt(1 - x ** 2)))
  acos (D x x')  = D (acos x) (- x' / (sqrt(1 - x ** 2)))
  atan (D x x')  = D (atan x) (x' / (1 + x ** 2))
  sinh (D x x')  = D (sinh x) (x' * cosh x)
  cosh (D x x')  = D (cosh x) (x' * sinh x)
  asinh (D x x') = D (asinh x) (x' / (sqrt(1 + x ** 2)))
  acosh (D x x') = D (acosh x) (x' / (sqrt(x ** 2 - 1)))
  atanh (D x x') = D (atanh x) (x' / (1 - x ** 2))


f :: Floating a => a -> a -> a -> a
f a b c = a * b ** 2 + sin c


-------------------------------------
-- 3.2 可微分プログラミング
-------------------------------------


type Model p x y =  forall z. Reifies z W
                 => BVar z p
                 -> BVar z x
                 -> BVar z y


linReg :: Model (Double, Double) Double Double
linReg (T2 a b) x = b * x + a


lossGrad :: (Backprop p, Backprop y, Num y)
         => Model p x y
         -> x           -- 入力データ
         -> y           -- 出力データ
         -> p           -- パラメータ
         -> p           -- パラメータによる勾配
lossGrad f x y = gradBP $ \p -> (f p (auto x) - auto y) ^ 2


trainModel :: (Fractional p, Backprop p, Num y, Backprop y)
           => Model p x y -- 学習モデル
           -> p           -- 初期パラメータ
           -> [(x,y)]     -- 観測データ
           -> p           -- 更新後のパラメータ
trainModel f = foldl $ \p (x,y) -> p - 0.1 * lossGrad f x y p


-------------------------------------
-- 4 ニューラルネットワークを実装する
-------------------------------------


logistic :: Floating a => a -> a
logistic x = 1 / (1 + exp (-x))


dense :: (KnownNat i, KnownNat o) => Model (L o i,  R o) (R i) (R o)
dense (T2 w b) x = w #> x + b


perceptron :: (KnownNat i, KnownNat o) => Model _ (R i) (R o)
perceptron p = logistic . dense p


-------------------------------------
-- 4.1 多層パーセプトロン
-------------------------------------


(<~) :: (Backprop p, Backprop q)
     => Model  p     b c
     -> Model     q  a b
     -> Model (p, q) a c
(f <~ g) (T2 p q) = f p . g q

infixr 8 <~


model :: (KnownNat i, KnownNat o) => Model _ (R i) (R o)
model = perceptron @4 <~ perceptron


-------------------------------------


main :: IO ()
main = do
  iW1 <- randn :: IO (L 1 4)
  iW2 <- randn :: IO (L 4 2)
  let iB1 = randomVector 0 Gaussian :: R 1
      iB2 = randomVector 0 Gaussian :: R 4
      samples = [(vec2 0 0, 0), (vec2 1 0, 1), (vec2 0 1, 1), (vec2 1 1, 0)]
      trained = trainModel model ((iW1, iB1), (iW2, iB2)) $ take 10000 (cycle samples)

  print $ evalBP2 model trained (vec2 0 0)
  print $ evalBP2 model trained (vec2 0 1)
  print $ evalBP2 model trained (vec2 1 0)
  print $ evalBP2 model trained (vec2 1 1)