My Flux model chain works perfectly when given an input or when calculating the loss of my training data. However, when I try to calculate the gradient my model chain undergoes unexpected behaviour.
My input data is a 3D tensor with shape (n x m x p). My model chain begins with c=Chain(DataLoader, rnns. This loads p 2D slices of the 3D tensor into the rnns function which for loops through doing [rnn(slice) in slices][end] (using a Flux RNN). The result of this operation is m vectors (columns) which have passed through the RNN. Thus it outputs an (n x m) matrix. This works for calculating output and loss.
However, this does not work when calculating the gradient. When calculating the gradient, for some reason the DataLoader simply converts the 3D CuArray to a 3D (cpu) Array. It does not batch the data into DataLoader type slices as it does when running input normally through the model chain or when calculating the loss.
I have attached a MWE below. As you can see, everything works as expected until the last line of code, the gradient calculation, which gives an error because the for loop through the rnn is iterating over the wrong type (Float32 Element instead of CuArrays in the DataLoader).
Thank you,
Jack
using Flux
using MLUtils
using CUDA
struct RNNModel
rnn_model::Flux.Recur
end
function (m::RNNModel)(data)
cur = nothing
for i in data # dataloader
println("hi!")
println(typeof(i))
println("typed")
println(i)
println("printed")
cur = m.rnn_model(i::CuArray{Float32, 3, CUDA.Mem.DeviceBuffer})
end
return cur
end
Flux.@functor RNNModel
n = 5
m = 4
p = 3
function printdata(data)
println(data)
println(typeof(data))
println("data printed")
return data
end
rnn = RNN(n=>n) |> gpu
rnns = RNNModel(rnn) |> gpu
c = Chain(DataLoader, printdata, rnns, softmax)
loss(x, y_target) = Flux.Losses.crossentropy(c(x), y_target)
x_train = randn(n, m, p) |> gpu
y_train = randn(n, m, 1)|>softmax |>gpu
c(x_train)
loss(x_train, y_train)
gradient(()->loss(x_train, y_train), Flux.params(c))
). The only reason we haven’t made it more prominent is because it’s unclear whether this should be the default API going forward. The main concern is how one would distinguish between a 3D array that represents a whole (batched) sequence, and one that represents a single timestep of a higher-dimensional batch (e.g. for a conv-rnn).