on master help sign in
Reimplement pooling backwards for speed
b914b91Huw Campbell
src/Grenade/Layers/Internal/Pooling.hs+31−7
@@ -5,6 +5,7 @@ module Grenade.Layers.Internal.Pooling (
5 , poolBackwardList
5 , poolBackwardList
6 ) where
6 ) where
7
7
8+import Data.Foldable ( forM_ )
8import Data.Function ( on )
9import Data.Function ( on )
9import Data.List ( maximumBy )
10import Data.List ( maximumBy )
10
11
@@ -12,7 +13,6 @@ import Numeric.LinearAlgebra hiding ( uniformSample, konst )
12import qualified Numeric.LinearAlgebra as LA
13import qualified Numeric.LinearAlgebra as LA
13import qualified Numeric.LinearAlgebra.Devel as U
14import qualified Numeric.LinearAlgebra.Devel as U
14
15
15−
16import Grenade.Layers.Internal.Convolution
16import Grenade.Layers.Internal.Convolution
17
17
18poolForward :: Int -> Int -> Int -> Int -> Int -> Int -> Matrix Double -> Matrix Double
18poolForward :: Int -> Int -> Int -> Int -> Int -> Int -> Matrix Double -> Matrix Double
@@ -42,15 +42,30 @@ poolBackwardList krows kcols srows scols inRows inCols inputMatrices =
42 let starts = fittingStarts inRows krows srows inCols kcols scols
42 let starts = fittingStarts inRows krows srows inCols kcols scols
43 in uncurry (poolBackwardFit starts krows kcols) <$> inputMatrices
43 in uncurry (poolBackwardFit starts krows kcols) <$> inputMatrices
44
44
45+-- poolBackwardFit :: [(Int,Int)] -> Int -> Int -> Matrix Double -> Matrix Double -> Matrix Double
46+-- poolBackwardFit starts krows kcols inputMatrix gradientMatrix =
47+-- let inRows = rows inputMatrix
48+-- inCols = cols inputMatrix
49+-- inds = fmap (\start -> unsafeMaxIndex $ subMatrix start (krows, kcols) inputMatrix) starts
50+-- grads = toList $ flatten gradientMatrix
51+-- grads' = zip3 starts grads inds
52+-- accums = fmap (\((stx',sty'),grad,(inx, iny)) -> ((stx' + inx, sty' + iny), grad)) grads'
53+-- in accum (LA.konst 0 (inRows, inCols)) (+) accums
54+
45poolBackwardFit :: [(Int,Int)] -> Int -> Int -> Matrix Double -> Matrix Double -> Matrix Double
55poolBackwardFit :: [(Int,Int)] -> Int -> Int -> Matrix Double -> Matrix Double -> Matrix Double
46−poolBackwardFit starts krows kcols inputMatrix gradientMatrix =
56+poolBackwardFit starts krows kcols inputMatrix gradientMatrix = U.runSTMatrix $ do
47 let inRows = rows inputMatrix
57 let inRows = rows inputMatrix
48 inCols = cols inputMatrix
58 inCols = cols inputMatrix
49− inds = fmap (\start -> unsafeMaxIndex $ subMatrix start (krows, kcols) inputMatrix) starts
59+ gradCol = cols gradientMatrix
50− grads = toList $ flatten gradientMatrix
60+ extent = (krows, kcols)
51− grads' = zip3 starts grads inds
61+
52− accums = fmap (\((stx',sty'),grad,(inx, iny)) -> ((stx' + inx, sty' + iny), grad)) grads'
53− in accum (LA.konst 0 (inRows, inCols)) (+) accums
62+ retM <- U.newMatrix 0 inRows inCols
63+
64+ forM_ (zip [0..] starts) $ \(ix, start) -> do
65+ let loc = unsafeMaxIndexSubMatrix start extent inputMatrix
66+ uncurry (unsafeModifyMatrix retM) loc ((+) $ uncurry (U.atM' gradientMatrix) $ divMod ix gradCol)
67+
68+ return retM
54
69
55unsafeMaxElement :: Matrix Double -> Double
70unsafeMaxElement :: Matrix Double -> Double
56unsafeMaxElement m = uncurry (U.atM' m) $ unsafeMaxIndex m
71unsafeMaxElement m = uncurry (U.atM' m) $ unsafeMaxIndex m
@@ -61,3 +76,12 @@ unsafeMaxIndex m =
61 mcols = [0 .. cols m - 1]
76 mcols = [0 .. cols m - 1]
62 pairs = concatMap ( \r -> fmap (\c -> (r , c)) mcols ) mrows
77 pairs = concatMap ( \r -> fmap (\c -> (r , c)) mcols ) mrows
63 in maximumBy (compare `on` uncurry (U.atM' m)) pairs
78 in maximumBy (compare `on` uncurry (U.atM' m)) pairs
79+
80+
81+unsafeMaxIndexSubMatrix :: (Int,Int) -> (Int,Int) -> Matrix Double -> (Int, Int)
82+unsafeMaxIndexSubMatrix (startRow, startCol) (extentRow, extentCold) m =
83+ let mrows = [startRow .. startRow + extentRow - 1]
84+ mcols = [startCol .. startCol + extentCold - 1]
85+ pairs = concatMap ( \r -> fmap (\c -> (r , c)) mcols ) mrows
86+ in maximumBy (compare `on` uncurry (U.atM' m)) pairs
87+