I'm using Julia Flux.jl to train models. But when I customize my model, there is an issue the variable reference in gradient pullback function doesn't behave as expected. For simplicity, the issue code is simplified and shown as following
let
x=rand(Float32, (10,1))
v_layer=[
Dense(10, 11),
Dense(10, 12),
]
ps=Flux.Params([
vcat((collect.(params.(v_layer)))...)...,
])
function model(x)
out=x
for (layer_idx, layer)=enumerate(v_layer)
@show layer_idx, size(x)
out=layer(x)
end
return out
end
gs=gradient(ps) do
sum(model(x))
end
end
If we run the code, the size of x in for loop will change according to the @show macro output as following
(layer_idx, size(x)) = (1, (10, 1))
(layer_idx, size(x)) = (2, (11, 1))
And we will get the following error message
DimensionMismatch("A has dimensions (12,10) but B has dimensions (11,1)")
I think the reference change of out in for loop should not change the reference of x, but that happened. Besides if we put the function model whole body inside gradient do as following
gs=gradient(ps) do
out=x
for (layer_idx, layer)=enumerate(v_layer)
@show layer_idx, size(x)
out=layer(x)
end
sum(out)
end
worked just fine. Why is this?
PS. A slightly more complex but more practical version code is as follows and has the same error.
let
v_layer_A=[
Dense(10, 11),
Dense(11, 12),
]
v_layer_B=[
Dense(10, 11),
Dense(10, 12),
]
x=rand(Float32, (10,1))
ps=Flux.Params([
vcat((collect.(params.(v_layer_A)))...)...,
vcat((collect.(params.(v_layer_B)))...)...,
])
function model(x)
out=x
for (idx, (layer_A, layer_B))=enumerate(zip(v_layer_A, v_layer_B))
@show idx, size(x)
out_A=layer_A(out)
out_B=layer_B(x)
out=out_A.*out_B
end
return out
end
gs=gradient(ps) do
sum(model(x))
end
end