I’m trying to do something similar to what Flux.Conv does. I’ve tried to use the Conv layer but it seems to be opinionated about how the array is structured and what it outputs and so doesn’t quite match with what I want. I have a 3 dim matrix of features x series x batch as input to the layer, and I want to run a Dense on flattened chunks of features x i:i+k x batch across that matrix.
It works on the CPU, but I want it to run on the GPU. Here’s a general idea of it:
mapreduce( i -> reshape(dense(Flux.flatten(selectdim(x, dim, i:i+plus))), sizeOut), (x1, x2) -> cat(x1, x2; dims=dim), 1:(size(x, dim)-plus) )
Flux.Conv is able to do the something similar on the GPU:
using Flux, CUDA CUDA.allowscalar(false) gpu(Flux.Conv((2,2), 4=>2))(gpu(rand(2,2,4,1)))
Here’s an MWE to show the error that selectdim causes:
using Flux, CUDA CUDA.allowscalar(false) den = Flux.Dense(30 => 4) |> gpu x = [i + j / 10 for i in 1:10, j in 1:20, b in 1:5] |> gpu x2 = Flux.flatten(selectdim(x, 2, 1:3)) den(x2) # ERROR: Scalar indexing is disallowed.
So, there must be a way to do it because Flux.Conv works. I tried to look at Conv’s source code, but it calls NNlib.conv and that is implemented with multiple stages of generated functions which is a bit difficult for me to piece together.
Any suggestions on how this might be accomplished?