@@ -5,6 +5,7 @@ module Grenade.Layers.Internal.Pooling (
55 , poolBackwardList
66 ) where
77
8+import Data.Foldable ( forM_ )
89import Data.Function ( on )
910import Data.List ( maximumBy )
1011
@@ -12,7 +13,6 @@ import Numeric.LinearAlgebra hiding ( uniformSample, konst )
1213import qualified Numeric.LinearAlgebra as LA
1314import qualified Numeric.LinearAlgebra.Devel as U
1415
15−
1616import Grenade.Layers.Internal.Convolution
1717
1818poolForward :: Int -> Int -> Int -> Int -> Int -> Int -> Matrix Double -> Matrix Double
@@ -42,15 +42,30 @@ poolBackwardList krows kcols srows scols inRows inCols inputMatrices =
4242 let starts = fittingStarts inRows krows srows inCols kcols scols
4343 in uncurry (poolBackwardFit starts krows kcols) <$> inputMatrices
4444
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+
4555poolBackwardFit :: [(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
4757 let inRows = rows inputMatrix
4858 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
59+ gradCol = cols gradientMatrix
60+ extent = (krows, kcols)
61+
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
5469
5570unsafeMaxElement :: Matrix Double -> Double
5671unsafeMaxElement m = uncurry (U.atM' m) $ unsafeMaxIndex m
@@ -61,3 +76,12 @@ unsafeMaxIndex m =
6176 mcols = [0 .. cols m - 1]
6277 pairs = concatMap ( \r -> fmap (\c -> (r , c)) mcols ) mrows
6378 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+