@@ -0,0 +1,238 @@
1+module Grenade.Layers.Convolution.Internal (
2+ im2col
3+-- , im2colUnsafe
4+ , vid2col
5+ , col2im
6+ , col2imFit
7+ , col2vid
8+
9+ , col2vidUnsafe
10+ , col2imUnsafe
11+ , im2colUnsafe
12+ , vid2colUnsafe
13+ , fittingStarts
14+ ) where
15+
16+import Control.Monad.ST
17+import Control.Parallel.Strategies ( parMap, rseq )
18+
19+import Data.STRef
20+import Data.Foldable ( forM_ )
21+
22+import Numeric.LinearAlgebra hiding ( uniformSample, konst )
23+import qualified Numeric.LinearAlgebra as LA
24+import qualified Numeric.LinearAlgebra.Devel as U
25+
26+im2col :: Int -> Int -> Int -> Int -> Matrix Double -> Matrix Double
27+im2col nrows ncols srows scols m =
28+ let starts = fittingStarts (rows m) nrows srows (cols m) ncols scols
29+ in im2colFit starts nrows ncols m
30+
31+im2colFit :: [(Int,Int)] -> Int -> Int -> Matrix Double -> Matrix Double
32+im2colFit starts nrows ncols m =
33+ let imRows = fmap (\start -> flatten $ subMatrix start (nrows, ncols) m) starts
34+ in fromRows imRows
35+
36+vid2col :: Int -> Int -> Int -> Int -> Int -> Int -> [Matrix Double] -> Matrix Double
37+vid2col nrows ncols srows scols inputrows inputcols ms =
38+ let starts = fittingStarts inputrows nrows srows inputcols ncols scols
39+ subs = parMap rseq (im2colFit starts nrows ncols) ms
40+ in foldl1 (|||) subs
41+
42+col2vid :: Int -> Int -> Int -> Int -> Int -> Int -> Matrix Double -> [Matrix Double]
43+col2vid krows kcols srows scols drows dcols m =
44+ let starts = fittingStart (cols m) (krows * kcols) (krows * kcols)
45+ r = rows m
46+ mats = fmap (\s -> subMatrix (0,s) (r, krows * kcols) m) starts
47+ colSts = fittingStarts drows krows srows dcols kcols scols
48+ in parMap rseq (col2imFit colSts krows kcols drows dcols) mats
49+
50+col2im :: Int -> Int -> Int -> Int -> Int -> Int -> Matrix Double -> Matrix Double
51+col2im krows kcols srows scols drows dcols m =
52+ let starts = fittingStarts drows krows srows dcols kcols scols
53+ in col2imFit starts krows kcols drows dcols m
54+
55+-- | These functions are not even remotely safe, but it's only called from the statically typed
56+-- commands, so we should be good ?!?!?
57+-- Returns the starting sub matrix locations which fit inside the larger matrix for the
58+-- convolution. Takes into account the stride and kernel size.
59+fittingStarts :: Int -> Int -> Int -> Int -> Int -> Int -> [(Int,Int)]
60+fittingStarts nrows kernelrows steprows ncols kernelcols stepcolsh =
61+ let rs = fittingStart nrows kernelrows steprows
62+ cs = fittingStart ncols kernelcols stepcolsh
63+ ls = sequence [rs, cs]
64+ in fmap (\[a,b] -> (a,b)) ls
65+
66+-- | Returns the starting sub vector which fit inside the larger vector for the
67+-- convolution. Takes into account the stride and kernel size.
68+fittingStart :: Int -> Int -> Int -> [Int]
69+fittingStart width kernel steps =
70+ let go left | left + kernel < width
71+ = left : go (left + steps)
72+ | left + kernel == width
73+ = [left]
74+ | otherwise
75+ = error "Kernel and step do not fit in matrix."
76+ in go 0
77+
78+col2imFit :: [(Int,Int)] -> Int -> Int -> Int -> Int -> Matrix Double -> Matrix Double
79+col2imFit starts krows kcols drows dcols m =
80+ let indicies = (\[a,b] -> (a,b)) <$> sequence [[0..(krows-1)], [0..(kcols-1)]]
81+ convs = fmap (zip indicies . toList) . toRows $ m
82+ pairs = zip convs starts
83+ accums = concatMap (\(conv',(stx',sty')) -> fmap (\((ix,iy), val) -> ((ix + stx', iy + sty'), val)) conv') pairs
84+ in accum (LA.konst 0 (drows, dcols)) (+) accums
85+
86+-- void col2im_cpu(const Dtype* data_col, const int channels,
87+-- const int height, const int width, const int kernel_h, const int kernel_w,
88+-- const int pad_h, const int pad_w,
89+-- const int stride_h, const int stride_w,
90+-- const int dilation_h, const int dilation_w,
91+-- Dtype* data_im) {
92+-- caffe_set(height * width * channels, Dtype(0), data_im);
93+-- const int output_h = (height + 2 * pad_h -
94+-- (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1;
95+-- const int output_w = (width + 2 * pad_w -
96+-- (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;
97+-- const int channel_size = height * width;
98+-- for (int channel = channels; channel--; data_im += channel_size) {
99+-- for (int kernel_row = 0; kernel_row < kernel_h; kernel_row++) {
100+-- for (int kernel_col = 0; kernel_col < kernel_w; kernel_col++) {
101+-- int input_row = -pad_h + kernel_row * dilation_h;
102+-- for (int output_rows = output_h; output_rows; output_rows--) {
103+-- if (!is_a_ge_zero_and_a_lt_b(input_row, height)) {
104+-- data_col += output_w;
105+-- } else {
106+-- int input_col = -pad_w + kernel_col * dilation_w;
107+-- for (int output_col = output_w; output_col; output_col--) {
108+-- if (is_a_ge_zero_and_a_lt_b(input_col, width)) {
109+-- data_im[input_row * width + input_col] += *data_col;
110+-- }
111+-- data_col++;
112+-- input_col += stride_w;
113+-- }
114+-- }
115+-- input_row += stride_h;
116+-- }
117+-- }
118+-- }
119+-- }
120+-- }
121+
122+
123+-- let starts = fittingStart (cols m) (krows * kcols) (krows * kcols)
124+-- r = rows m
125+-- mats = fmap (\s -> subMatrix (0,s) (r, krows * kcols) m) starts
126+-- in parMap rseq (col2imUnsafe krows kcols srows scols drows dcols) mats
127+
128+col2imUnsafe :: Int -> Int -> Int -> Int -> Int -> Int -> Matrix Double -> Matrix Double
129+col2imUnsafe kernelRows kernelColumns strideRows strideColumns destinationRows destinationCols columnMatrix = U.runSTMatrix $ do
130+ let columnMatrixRows = rows columnMatrix
131+
132+ dataIm <- U.newMatrix 0 destinationRows destinationCols
133+
134+ offsetR <- newSTRef 0
135+ offsetC <- newSTRef 0
136+
137+ forM_ [0 .. columnMatrixRows - 1] $ \ir -> do
138+ inputColumn <- newSTRef 0
139+ forM_ [0 .. kernelRows -1] $ \kr ->
140+ forM_ [0 .. kernelColumns -1] $ \kc -> do
141+ ic <- readSTRef inputColumn
142+ offsetR' <- readSTRef offsetR
143+ offsetC' <- readSTRef offsetC
144+ U.modifyMatrix dataIm (kr + offsetR') (kc + offsetC') (+ atIndex columnMatrix (ir,ic))
145+ modifySTRef inputColumn (+1)
146+
147+ offsetC' <- readSTRef offsetC
148+ if offsetC' + kernelColumns < destinationCols
149+ then modifySTRef offsetC (+ strideColumns)
150+ else writeSTRef offsetC 0 >> modifySTRef offsetR (+ strideRows)
151+
152+ return dataIm
153+
154+col2vidUnsafe :: Int -> Int -> Int -> Int -> Int -> Int -> Matrix Double -> [Matrix Double]
155+col2vidUnsafe kernelRows kernelColumns strideRows strideColumns destinationRows destinationCols columnMatrix = runST $ do
156+ let columnMatrixRows = rows columnMatrix
157+ let filters = cols columnMatrix `div` (kernelRows * kernelColumns)
158+
159+ dataIms <- traverse (\_ -> U.newMatrix 0 destinationRows destinationCols) [0 .. filters-1]
160+
161+ offsetR <- newSTRef 0
162+ offsetC <- newSTRef 0
163+ offsetM <- newSTRef 0
164+
165+ forM_ dataIms $ \dataIm -> do
166+ offsetM' <- readSTRef offsetM
167+ forM_ [0 .. columnMatrixRows - 1] $ \ir -> do
168+ inputColumn <- newSTRef 0
169+ forM_ [0 .. kernelRows -1] $ \kr ->
170+ forM_ [0 .. kernelColumns -1] $ \kc -> do
171+ ic <- readSTRef inputColumn
172+ offsetR' <- readSTRef offsetR
173+ offsetC' <- readSTRef offsetC
174+ U.modifyMatrix dataIm (kr + offsetR') (kc + offsetC') (+ atIndex columnMatrix (ir, ic + offsetM'))
175+ modifySTRef inputColumn (+1)
176+
177+ offsetC' <- readSTRef offsetC
178+ if offsetC' + kernelColumns < destinationCols
179+ then modifySTRef offsetC (+ strideColumns)
180+ else writeSTRef offsetC 0 >> modifySTRef offsetR (+ strideRows)
181+
182+ writeSTRef offsetR 0
183+ writeSTRef offsetC 0
184+ modifySTRef offsetM (+ (kernelRows * kernelColumns))
185+
186+ traverse U.freezeMatrix dataIms
187+
188+vid2colUnsafe :: Int -> Int -> Int -> Int -> Int -> Int -> Int -> [Matrix Double] -> Matrix Double
189+vid2colUnsafe channels kernelRows kernelColumns striderows stridecols vidrows vidcols dataVid = U.runSTMatrix $ do
190+ let starts = fittingStarts vidrows kernelRows striderows vidcols kernelColumns stridecols
191+ matWidth = kernelRows * kernelColumns
192+ destinationRows = 1 + (vidrows - kernelRows) `div` striderows
193+ destinationCols = 1 + (vidcols - kernelColumns) `div` stridecols
194+ destinationSize = destinationRows * destinationCols
195+
196+ dataCol <- U.newMatrix 0 destinationSize (channels * matWidth)
197+
198+ offsetC <- newSTRef 0
199+
200+ forM_ dataVid $ \dataIm -> do
201+ inputRow <- newSTRef 0
202+ offsetC' <- readSTRef offsetC
203+ forM_ starts $ \(startRow, startCol) -> do
204+ inputColumn <- newSTRef 0
205+ inputRow' <- readSTRef inputRow
206+ forM_ [0 .. kernelRows -1] $ \kr ->
207+ forM_ [0 .. kernelColumns -1] $ \kc -> do
208+ inputColumn' <- readSTRef inputColumn
209+ U.modifyMatrix dataCol inputRow' (inputColumn' + offsetC') (+ atIndex dataIm (kr + startRow, kc + startCol))
210+ modifySTRef inputColumn (+1)
211+ modifySTRef inputRow (+1)
212+
213+ modifySTRef offsetC (+ matWidth)
214+
215+ return dataCol
216+
217+im2colUnsafe :: Int -> Int -> Int -> Int -> Int -> Int -> Matrix Double -> Matrix Double
218+im2colUnsafe kernelRows kernelColumns striderows stridecols vidrows vidcols dataIm = U.runSTMatrix $ do
219+ let starts = fittingStarts vidrows kernelRows striderows vidcols kernelColumns stridecols
220+ matWidth = kernelRows * kernelColumns
221+ destinationRows = 1 + (vidrows - kernelRows) `div` striderows
222+ destinationCols = 1 + (vidcols - kernelColumns) `div` stridecols
223+ destinationSize = destinationRows * destinationCols
224+
225+ dataCol <- U.newMatrix 0 destinationSize matWidth
226+
227+ inputRow <- newSTRef 0
228+ forM_ starts $ \(startRow, startCol) -> do
229+ inputColumn <- newSTRef 0
230+ inputRow' <- readSTRef inputRow
231+ forM_ [0 .. kernelRows -1] $ \kr ->
232+ forM_ [0 .. kernelColumns -1] $ \kc -> do
233+ inputColumn' <- readSTRef inputColumn
234+ U.modifyMatrix dataCol inputRow' inputColumn' (+ atIndex dataIm (kr + startRow, kc + startCol))
235+ modifySTRef inputColumn (+1)
236+ modifySTRef inputRow (+1)
237+
238+ return dataCol