@@ -1,4 +1,5 @@
11{-# LANGUAGE DataKinds #-}
2+{-# LANGUAGE BangPatterns #-}
23{-# LANGUAGE GADTs #-}
34{-# LANGUAGE ScopedTypeVariables #-}
45{-# LANGUAGE TypeOperators #-}
@@ -7,40 +8,76 @@
78{-# LANGUAGE FlexibleContexts #-}
89{-# LANGUAGE FlexibleInstances #-}
910module Grenade.Layers.Fuse (
10− Fuse (..)
11+ SubNetwork (..)
12+ , BranchNetwork (..)
1113 ) where
1214
15+import Data.Singletons.Prelude
16+
1317import Grenade.Core.Network
1418import Grenade.Core.Shape
19+import Grenade.Core.Runner
1520
1621-- | Fuse two layers into one layer.
1722-- This can be used to simplify a network if a complicated repeated structure is used.
18−data Fuse :: * -> * -> Shape -> Shape -> Shape -> * where
19− (:$$) :: (Layer x i h, Layer y h o)
20− => !x
21− -> !y
22− -> Fuse x y i h o
23−infixr 5 :$$
24−
25−instance (Show x, Show y) => Show (Fuse x y i h o) where
26− show (x :$$ y) = "(" ++ show x ++ " :$$ " ++ show y ++ ")"
27−
28−instance (Layer x i h, Layer y h o) => UpdateLayer (Fuse x y i h o) where
29− type Gradient (Fuse x y i h o) = (Gradient x, Gradient y)
30− runUpdate lr (x :$$ y) (x', y') =
31− let newX = runUpdate lr x x'
32− newY = runUpdate lr y y'
33− in (newX :$$ newY)
34− createRandom = (:$$) <$> createRandom <*> createRandom
35−
36−instance (Layer x i h, Layer y h o) => Layer (Fuse x y i h o) i o where
37− type Tape (Fuse x y i h o) i o = (Tape x i h, Tape y h o)
38− runForwards (x :$$ y) input =
39− let (xTape, yInput :: S h) = runForwards x input
40− (yTape, output :: S o) = runForwards y yInput
41− in ((xTape, yTape), output)
42−
43− runBackwards (x :$$ y) (xTape, yTape) backGradient =
44− let (y', yGrad :: S h) = runBackwards y yTape backGradient
45− (x', xGrad) = runBackwards x xTape yGrad
46− in ((x', y'), xGrad)
23+data SubNetwork :: [*] -> [Shape] -> * where
24+ SubNetwork :: Network sublayers subshapes -> SubNetwork sublayers subshapes
25+
26+instance CreatableNetwork sublayers subshapes => UpdateLayer (SubNetwork sublayers subshapes) where
27+ type Gradient (SubNetwork sublayers subshapes) = Gradients sublayers
28+ runUpdate lr (SubNetwork net) = SubNetwork . applyUpdate lr net
29+ createRandom = SubNetwork <$> randomNetwork
30+
31+instance (CreatableNetwork sublayers subshapes, i ~ (Head subshapes), o ~ (Last subshapes)) => Layer (SubNetwork sublayers subshapes) i o where
32+ type Tape (SubNetwork sublayers subshapes) i o = Tapes sublayers subshapes
33+
34+ runForwards (SubNetwork net) i =
35+ go i net
36+ where
37+ go :: forall js ss. (Last js ~ Last subshapes)
38+ => S (Head js) -- ^ input vector
39+ -> Network ss js -- ^ network to train
40+ -> (Tapes ss js, S (Last js))
41+ go !x (layer :~> n) =
42+ let (tape, forward) = runForwards layer x
43+ (tapes, answer) = go forward n
44+ in (tape :\> tapes, answer)
45+ go !x NNil
46+ = (TNil, x)
47+
48+ runBackwards (SubNetwork net) tapes o =
49+ go net tapes
50+ where
51+ go :: forall js ss. (Last js ~ Last subshapes)
52+ => Network ss js -- ^ network to train
53+ -> Tapes ss js -- ^ network to train
54+ -> (Gradients ss, S (Head js))
55+ go (layer :~> n) (tape :\> nt) =
56+ let (gradients, feed) = go n nt
57+ (layer', backGrad) = runBackwards layer tape feed
58+ in (layer' :/> gradients, backGrad)
59+ go NNil TNil
60+ = (GNil, o)
61+
62+
63+-- | Run two layers in parallel, combining their outputs. The way in which the output is combined is dependent
64+-- on the inut
65+data BranchNetwork :: * -> * -> * where
66+ BranchNetwork :: x -> y -> BranchNetwork x y
67+
68+instance (UpdateLayer x, UpdateLayer y) => UpdateLayer (BranchNetwork x y) where
69+ type Gradient (BranchNetwork x y) = (Gradient x, Gradient y)
70+ runUpdate lr (BranchNetwork x y) (x', y') = BranchNetwork (runUpdate lr x x') (runUpdate lr y y')
71+ createRandom = BranchNetwork <$> createRandom <*> createRandom
72+
73+instance (SingI i, SingI o, Layer x i o, Layer y i o) => Layer (BranchNetwork x y) i o where
74+ type Tape (BranchNetwork x y) i o = (Tape x i o, Tape y i o)
75+ runForwards (BranchNetwork x y) input =
76+ let (xT, xOut) = runForwards x input
77+ (yT, yOut) = runForwards y input
78+ in ((xT, yT), xOut + yOut)
79+
80+ runBackwards (BranchNetwork x y) (xTape, yTape) o =
81+ let (x', xB) = runBackwards x xTape o
82+ (y', yB) = runBackwards y yTape o
83+ in ((x', y'), xB + yB)