hasktorch-0.0.1.0: tests/Torch/Static/NN/LinearSpec.hs
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE CPP #-}
#if MIN_VERSION_base(4,12,0)
{-# LANGUAGE NoStarIsType #-}
#endif
{-# OPTIONS_GHC -fno-cse #-}
module Torch.Static.NN.LinearSpec where
import GHC.TypeLits
import Control.Monad (join, void)
import Data.Function ((&))
import Data.Foldable
import Debug.Trace
import Data.Maybe
import GHC.Generics (Generic)
import Test.Hspec
import Lens.Micro.Platform
import Numeric.Backprop
import System.IO.Unsafe
import Data.Generics.Product
import qualified Numeric.Backprop as B
import Debug.Trace
import Torch.Double as Torch
import Torch.Double.NN.Linear
import qualified Torch.Long as Long
reasonablyUnsafeVector :: (KnownDim m, KnownNat m) => [HsReal] -> Tensor '[m]
reasonablyUnsafeVector = unsafePerformIO . unsafeVector
{-# NOINLINE reasonablyUnsafeVector #-}
reasonablyUnsafeLongVector :: (KnownDim m, KnownNat m) => [Long.HsReal] -> Long.Tensor '[m]
reasonablyUnsafeLongVector = unsafePerformIO . Long.unsafeVector
{-# NOINLINE reasonablyUnsafeLongVector #-}
reasonablyUnsafeMatrix
:: All KnownDim '[m, n, n*m]
=> All KnownNat '[m, n, n*m]
=> [[HsReal]]
-> Tensor '[n,m]
reasonablyUnsafeMatrix = unsafePerformIO . unsafeMatrix
{-# NOINLINE reasonablyUnsafeMatrix #-}
xavier :: forall d . Dimensions d => IO (Tensor d)
xavier = case (fromIntegral <$> listDims (dims :: Dims d)) of
[] -> pure empty
a:_ -> pure $ constant (1 / realToFrac (fromIntegral a))
data FF2Network i h o = FF2Network
{ layer1 :: Linear i h
, layer2 :: Linear h o
} deriving (Generic, Show)
instance (KnownDim i, KnownDim h, KnownDim o) => Pairwise (FF2Network i h o) HsReal where
(FF2Network l0 l1) ^+ v = FF2Network (l0 ^+ v) (l1 ^+ v)
(FF2Network l0 l1) ^- v = FF2Network (l0 ^+ v) (l1 ^+ v)
(FF2Network l0 l1) ^* v = FF2Network (l0 ^+ v) (l1 ^+ v)
(FF2Network l0 l1) ^/ v = FF2Network (l0 ^+ v) (l1 ^+ v)
weightsL :: Lens' (Linear i o) (Tensor '[i, o])
weightsL = field @"getTensors" . _1
biasL :: Lens' (Linear i o) (Tensor '[o])
biasL = field @"getTensors" . _2
specupdate
:: forall i h o
. All KnownDim '[i, h, o]
=> FF2Network i h o
-> FF2Network i h o
-> FF2Network i h o
specupdate i g = FF2Network
{ layer1 = B.add (layer1 i) (layer1 g)
, layer2 = B.add (layer2 i) (layer2 g)
}
instance (KnownDim i, KnownDim h, KnownDim o) => Backprop (FF2Network i h o)
-- ========================================================================= --
main :: IO ()
main = hspec spec
spec :: Spec
spec = do
describe "a single linear layer" singleLayer
describe "a two-layer feed forward network" $ do
describe "with xavier initialization" twoLayerXavier
describe "forcing ReLU activity" twoLayerForceReLU
describe "overfitting to [0, 1]" $ do
describe "with one layer and binary cross-entropy" oneLayerOverfit
describe "with two layers and binary cross-entropy" twoLayerOverfit
-- describe "with logSoftmax and multiclass log-loss" $
-- twoLayerOverfit logSoftMax (classNLLCriterion (reasonablyUnsafeLongVector [1])) Torch.exp
-- ========================================================================= --
singleLayer :: Spec
singleLayer = do
ll :: Linear 3 2 <- runIO $ mkLinear xavier
xavierPurityCheck ll $ do
describe "the forward pass" $ do
let y = constant 5 :: Tensor '[3]
o = evalBP2 (linear) ll y
it "performs matrix multipication as you would expect" $ do
o =##= ((5/3)*3) + 1/2
o `elementsSatisfy` ((== 2) . length)
describe "the backward pass" $ do
let y = constant 1 :: Tensor '[3]
(_, (ll', o)) = backprop2 linear ll y
it "returns plain gradient of weights" $ weights ll' =##= 1 -- 1/2
it "returns plain gradient of bias" $ bias ll' =##= 3/2 -- 2/3
it "returns plain gradient of output tensor" $ o =##= 2/3 -- 1/3
-- ========================================================================= --
mkXavierNetwork :: All KnownDim '[i,h,o] => IO (FF2Network i h o)
mkXavierNetwork =
FF2Network
<$> mkLinear xavier
<*> mkLinear xavier
mkUniform :: All KnownDim '[i,h,o] => IO (FF2Network i h o)
mkUniform = do
g <- newRNG
manualSeed g 1
let Just rg = ord2Tuple (-1, 1)
w0 <- uniform g rg
w1 <- uniform g rg
pure $ FF2Network
(Linear (w0, constant 1))
(Linear (w1, constant 1))
mkGaussianNetwork :: All KnownDim '[i,h,o] => IO (FF2Network i h o)
mkGaussianNetwork = do
g <- newRNG
manualSeed g 1
let Just std = positive 2
FF2Network
<$> mkLinear (normal g 0 std)
<*> mkLinear (normal g 0 std)
ff2network
:: forall s i h o
. Reifies s W
=> All KnownDim '[i,h,o]
=> (forall s . Reifies s W => BVar s (Tensor '[o]) -> BVar s (Tensor '[o]))
-> Double
-> BVar s (FF2Network i h o) -- ^ ff2network architecture
-> BVar s (Tensor '[i]) -- ^ input
-> BVar s (Tensor '[1, o]) -- ^ output
ff2network final lr arch inp
= linear {-lr-} (arch ^^. field @"layer1") inp
& relu
& linear {-lr-} (arch ^^. field @"layer2")
& final
& foo
where
foo t = unsqueeze1dBP (dim :: Dim 0) t
twoLayerXavier :: Spec
twoLayerXavier = do
ll :: FF2Network 4 6 2 <- runIO mkXavierNetwork
describe "the forward pass" $ do
describe "with xavier instantiation" $ do
xavierPurityCheck (ll ^. field @"layer1") $
xavierPurityCheck (ll ^. field @"layer2") $
describe "with input all positive input" $ do
let y = constant 4 :: Tensor '[4]
(o, _) = backprop2 (ff2network softmax undefined) ll y
it "performs matrix multipication as you would expect" $ o `approx` [1/2, 1/2]
describe "with input that drops all values via ReLU" $ do
let y = constant (-1) :: Tensor '[4]
(o, _) = backprop2 (ff2network softmax undefined) ll y
it "performs matrix multipication as you would expect" $ o `approx` [1/2, 1/2]
where
approx = lapproximately 0.0001
twoLayerForceReLU :: Spec
twoLayerForceReLU = do
describe "operations that force relu activity" $ do
let o1 = evalBP2 (relu .: linear {-1-}) (ff2 ^. field @"layer1") oneInput
o2 = evalBP2 ( linear {-1-}) (ff2 ^. field @"layer2") o1
gin :: Tensor '[4]
(out, (gff2, gin)) = backprop2 (ff2network logSoftMax 1) ff2 oneInput
describe "dropping half the gradient during ReLU" $ do
describe "the forward pass" $ do
it "returns [0,0,0,4,4,4] after the first layer" $
tensordata o1 `shouldBe` [0,0,0,4,4,4]
it "returns [-12,0] after the second layer" $
tensordata o2 `shouldBe` [-12, 0]
it "returns [0, 1] as the output" $ do
Torch.exp out `lapprox` [0, 1]
describe "the backward pass" $ do
it "returns a half zero-d out layer 1 gradient" $ do
(gff2 ^. field @"layer1" . weightsL) `approx` l1weightgrad
it "returns a quarter zero-d out layer 2 gradient" $ do
(gff2 ^. field @"layer2" . weightsL) `approx` l2weightgrad
it "returns a [3,3,3] input gradient" $ do
gin `lapprox` replicate 4 (-3)
where
eps = 0.0001
approx :: Tensor d -> Tensor d -> IO ()
approx = approximately eps
lapprox :: Tensor d -> [HsReal] -> IO ()
lapprox = lapproximately eps
ff2 :: FF2Network 4 6 2
ff2 = FF2Network
(Linear (reasonablyUnsafeMatrix $ replicate 4 [ -1, -1, -1, 1, 1, 1], constant 0))
(Linear (reasonablyUnsafeMatrix $ replicate 6 [-1, 0], constant 0))
oneInput :: Tensor '[4]
oneInput = constant 1
l1weightgrad :: (Tensor '[4, 6])
l1weightgrad = reasonablyUnsafeMatrix $ replicate 4 [ 0, 0, 0,-1,-1,-1]
l2weightgrad :: (Tensor '[6, 2])
l2weightgrad = reasonablyUnsafeMatrix $ replicateN 3
[ [ 0, 0]
, [ 4,-4]
]
twoLayerOverfit :: Spec
twoLayerOverfit = do
net0 <- runIO $ do
g <- newRNG
manualSeed g 1
l0 <- (Linear . (,constant 1)) <$> uniform g rg
l1 <- (Linear . (,constant 1)) <$> uniform g rg
pure (l0, l1)
it "returns around 50-50 on uniform random initialization" . void $ do
let [l0, r0] = tensordata $ infer net0
let pointapprox pred truth = Prelude.abs (pred - truth) < 0.3
(l0, r0) `shouldSatisfy` (\(px, py) -> pointapprox px 0.5 && pointapprox py 0.5)
it "backprops to yield a loss smaller than its prior" . void $ do
let [l0, r0] = tensordata $ infer net0
let lr = (-0.001) :: HsReal
let (o, _) = bprop net0
(fnet, (fo, fl, fr)) <-
foldlM (\(net, (o, l, r)) i -> do
let (o, (Linear (gw0, gb0), Linear (gw1, gb1))) = bprop net
let net' = B.add net (Linear (gw0 ^* lr, gb0 ^* lr), Linear (gw1 ^* lr, gb1 ^* lr))
let (o', grad') = bprop net'
let [l', r'] = tensordata $ infer net'
o `shouldSatisfy` (> o')
pure (net', (o', l', r'))
) (net0, (o, l0, r0)) [1..100]
let pointapprox pred truth = Prelude.abs (pred - truth) < 0.01
(fl, fr) `shouldSatisfy` (\(px, py) -> pointapprox px 0.0 && pointapprox py 1.0)
where
Just rg = ord2Tuple (-1, 1)
x :: Tensor '[4]
x = constant 1
answer :: Tensor '[2]
answer = reasonablyUnsafeVector [0,1]
arch :: Reifies s W => BVar s (Linear 4 6, Linear 6 2) -> BVar s (Tensor '[4]) -> BVar s (Tensor '[2])
arch arch inp
= linear (arch ^^. _1) inp
& relu
& linear (arch ^^. _2)
& softmax
infer :: (Linear 4 6, Linear 6 2) -> Tensor '[2]
infer net = evalBP2 arch net x
bprop net = (fromJust $ get1d o 0, g)
where
(o, (g, _)) = backprop2 (bCECriterion answer .: arch) net x
oneLayerOverfit :: Spec
oneLayerOverfit = do
net0 <- runIO (newRNG >>= \g -> manualSeed g 1 >> (Linear . (,constant 1)) <$> uniform g rg)
it "returns around 50-50 on uniform random initialization" . void $ do
let [l0, r0] = tensordata $ infer net0
let pointapprox pred truth = Prelude.abs (pred - truth) < 0.3
(l0, r0) `shouldSatisfy` (\(px, py) -> pointapprox px 0.5 && pointapprox py 0.5)
it "backprops to yield a loss smaller than its prior" . void $ do
let [l0, r0] = tensordata $ infer net0
let lr = (-0.1) :: HsReal
let (o, _) = bprop net0
(fnet, (fo, fl, fr)) <-
foldlM (\(net, (o, l, r)) i -> do
let (o, Linear (gw, gb)) = bprop net
let net' = B.add net (Linear (gw ^* lr, gb ^* lr))
let (o', grad') = bprop net'
let [l', r'] = tensordata $ infer net'
o `shouldSatisfy` (> o')
pure (net', (o', l', r'))
) (net0, (o, l0, r0)) [1..100]
let pointapprox pred truth = Prelude.abs (pred - truth) < 0.01
(fl, fr) `shouldSatisfy` (\(px, py) -> pointapprox px 0.0 && pointapprox py 1.0)
where
Just rg = ord2Tuple (-1, 1)
x :: Tensor '[6]
x = constant 1
answer :: Tensor '[2]
answer = (reasonablyUnsafeVector [0,1])
arch :: Reifies s W => BVar s (Linear 6 2) -> BVar s (Tensor '[6]) -> BVar s (Tensor '[2])
arch a b = softmax $ linear a b
infer :: Linear 6 2 -> Tensor '[2]
infer net = evalBP2 arch net x
bprop net = (fromJust $ get1d o 0, g)
where
(o, (g, _)) = backprop2 (bCECriterion answer .: arch) net x
xavierPurityCheck :: forall i o . (KnownDim i, KnownDim o) => Linear i o -> Spec -> Spec
xavierPurityCheck ll tests =
it (header ++ "initializes with xavier correctly")
( weights ll =##= 1/i
>> bias ll =##= 1/o)
>> tests
>> it (header ++ "leaves weights unchanged") (weights ll =##= 1/i)
>> it (header ++ "leaves bias unchanged") (bias ll =##= 1/o)
where
header :: String
header = "[ref-check] " ++ unwords ["Linear", show (truncate i), show (truncate o)] ++ ": "
i, o :: Double
i = fromIntegral (dimVal (dim :: Dim i))
o = fromIntegral (dimVal (dim :: Dim o))
_lapproximately :: ([Double] -> Bool) -> Double -> Tensor d -> [Double] -> IO ()
_lapproximately pred e o dist = let os = tensordata o in
zipWith (Prelude.abs .: subtract) os dist `shouldSatisfy` pred
_approximately :: ([Double] -> Bool) -> Double -> Tensor d -> Tensor d -> IO ()
_approximately pred e o dist = _lapproximately pred e o (tensordata dist)
approximately e = _approximately (all (< e)) e
notCloseTo e = _approximately (all (> e)) e
lapproximately e = _lapproximately (all (< e)) e
lnotCloseTo e = _lapproximately (all (> e)) e
elementsSatisfy :: Tensor d -> ([Double] -> Bool) -> IO ()
elementsSatisfy o pred = tensordata o `shouldSatisfy` pred
replicateN :: Int -> [a] -> [a]
replicateN inner = concatMap (replicate inner)
(=##=) :: Tensor d -> Double -> IO ()
(=##=) o v = elementsSatisfy o (all (== v))
infixl 2 =##=