# Flux.jl Inconsistent Training on Custom Architecture
**URL:** https://discourse.julialang.org/t/flux-jl-inconsistent-training-on-custom-architecture/63326
**Category:** New to Julia
**Tags:** question, flux, autodiff
**Created:** [June 21, 2021, 7:55pm UTC](https://discourse.julialang.org/t/flux-jl-inconsistent-training-on-custom-architecture/63326 "2021-06-21T19:55:45Z")
**Posts on this page:** 10
**Page:** 1
### Author:  [@than123](https://discourse.julialang.org/u/than123)
#### Post date: [June 21, 2021, 7:55pm UTC](https://discourse.julialang.org/t/flux-jl-inconsistent-training-on-custom-architecture/63326/1 "2021-06-21T19:55:46Z")
I am currently implementing a toy version of the architecture as found in:
[Deep learning for universal linear embeddings of nonlinear dynamics | Nature Communications](https://www.nature.com/articles/s41467-018-07210-0), which is effectively an encoder → affine → decoder process.
I have already implemented the architecture and have been testing on single trajectory datasets. This is the struct I am employing as the model:
```julia
struct A
encoder
decoder
K
end
```
where the encoder and decoder are MLPs and K is simply a matrix. In order to move my testing on to multiple trajectories, I am employing this model:
```julia
struct B
encoder
decoder
K::Vector
end
```
In this case, K in this model holds one of the K’s in the above model per trajectory. So, these two architectures should be equivalent if training on a single trajectory dataset.
However, I am not getting this expected behavior. For a single trajectory dataset, `A` fits well while `B` does not converge whatsoever no matter the training time or parameters used. In fact, the loss almost stays constant throughout training but jitters somewhat.
While attempting to print the norm of the gradient for `B`, I get an error which suggest to me that the gradient is empty. Here is an MWE:
```julia
using Flux
using LinearAlgebra
Xs = [rand(2,10)]
mutable struct B
mlp
K::Vector
end
function loss(Xs,model)
l = 0
for i in 1:length(Xs)
l += norm(model.K[i])
end
return l
end
model = B(Dense(2,2),[rand(2,2)])
ps = params(model.mlp,model.K)
gs = gradient(() -> loss(Xs,model),ps)
norm(gs)
```
which gives the following error:
```julia
**ERROR:** MethodError: no method matching iterate(::Nothing)
Closest candidates are:
iterate(::DataStructures.TrieIterator) at /Users/tylerhan/.julia/packages/DataStructures/ixwFs/src/trie.jl:112
iterate(::DataStructures.TrieIterator, ::Any) at /Users/tylerhan/.julia/packages/DataStructures/ixwFs/src/trie.jl:112
iterate(::Cmd) at process.jl:638
...
Stacktrace:
[1] **isempty(** ::Nothing **)** at **./essentials.jl:737**
[2] **norm(** ::Nothing, ::Int64 **)** at **/Users/julia/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.5/LinearAlgebra/src/generic.jl:605** (repeats 2 times)
[3] **generic_normInf(** ::Zygote.Grads **)** at **/Users/julia/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.5/LinearAlgebra/src/generic.jl:446**
[4] **normInf** at **/Users/julia/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.5/LinearAlgebra/src/generic.jl:536** [inlined]
[5] **generic_norm2(** ::Zygote.Grads **)** at **/Users/julia/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.5/LinearAlgebra/src/generic.jl:477**
[6] **norm2** at **/Users/julia/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.5/LinearAlgebra/src/generic.jl:538** [inlined]
[7] **norm(** ::Zygote.Grads, ::Int64 **)** at **/Users/julia/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.5/LinearAlgebra/src/generic.jl:607**
[8] **norm(** ::Zygote.Grads **)** at **/Users/julia/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.5/LinearAlgebra/src/generic.jl:605**
[9] top-level scope at **REPL[9]:1**
```
I am not certain of this, but perhaps this means that the second model’s gradient is empty? I do not get such errors for the first model.
I’d like to know if there is a fix for such unexpected behavior or perhaps there is a better way to implement the second model (I have looked at Flux’s “advanced model building” section but it was not obvious to me how it might help).
I am currently using Flux v0.12.1.
---
### Author:  [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)
#### Post date: [June 21, 2021, 10:14pm UTC](https://discourse.julialang.org/t/flux-jl-inconsistent-training-on-custom-architecture/63326/2 "2021-06-21T22:14:45Z")
Please see [Please read: make it easier to help you](https://discourse.julialang.org/t/please-read-make-it-easier-to-help-you/14757). We can try to work out the issue by speculating about it, but _the_ best way to get a resolution is to post a MWE that demonstrates it. For Flux stuff, that includes dummy data, library versions used and a full stacktrace of any relevant errors.
---
### Author:  [@than123](https://discourse.julialang.org/u/than123)
#### Post date: [June 22, 2021, 1:47am UTC](https://discourse.julialang.org/t/flux-jl-inconsistent-training-on-custom-architecture/63326/3 "2021-06-22T01:47:50Z")
I see! Will do. Thanks for letting me know.
---
### Author:  [@than123](https://discourse.julialang.org/u/than123)
#### Post date: [June 22, 2021, 2:16am UTC](https://discourse.julialang.org/t/flux-jl-inconsistent-training-on-custom-architecture/63326/4 "2021-06-22T02:16:58Z")
The description has been updated! Please let me know if there are other things I could add which might help.
---
### Author:  [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)
#### Post date: [June 22, 2021, 2:38am UTC](https://discourse.julialang.org/t/flux-jl-inconsistent-training-on-custom-architecture/63326/5 "2021-06-22T02:38:49Z")
Can you also provide an end-to-end, minimal running example that throws that error (i.e. a MWE)? As-is there’s nothing inherent in the definition of either struct that would cause an issue, so having that full context is important.
---
### Author:  [@than123](https://discourse.julialang.org/u/than123)
#### Post date: [June 22, 2021, 3:51am UTC](https://discourse.julialang.org/t/flux-jl-inconsistent-training-on-custom-architecture/63326/6 "2021-06-22T03:51:22Z")
Edited. Thank you for bearing with me!
---
### Author:  [@than123](https://discourse.julialang.org/u/than123)
#### Post date: [June 23, 2021, 6:43pm UTC](https://discourse.julialang.org/t/flux-jl-inconsistent-training-on-custom-architecture/63326/7 "2021-06-23T18:43:12Z")
I have found a workaround by simply augmenting K, but for the future’s sake, I would still like to know what exactly here is not permissible.
---
### Author:  [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)
#### Post date: [June 23, 2021, 6:54pm UTC](https://discourse.julialang.org/t/flux-jl-inconsistent-training-on-custom-architecture/63326/8 "2021-06-23T18:54:25Z")
`gradient` with implicit params (i.e. what `params` returns) will return a `Grads` struct. This is not an array, but more a bag of arrays. To calculate the `norm` of that, you’d have to iterate through it and calculate the norm of each element. Note that this also requires checking for `nothing`, as that’s the value used for parameters that aren’t involved in the gradient calculation.
As for your workaround, I’m not sure what “augmenting K” means. Could you post an updated example with this augmentation that works as you’d expect?
---
### Author:  [@than123](https://discourse.julialang.org/u/than123)
#### Post date: [June 23, 2021, 7:31pm UTC](https://discourse.julialang.org/t/flux-jl-inconsistent-training-on-custom-architecture/63326/9 "2021-06-23T19:31:36Z")
I see, but I don’t get an error involving `nothing` for the first model. And in the second model, all the parameters should be used for the gradient.
By “augmenting K”, I just mean simply stacking the corresponding K matrices, as in:
```julia
newK = vcat(oldK...)
```
and I identify the relevant K matrices in the loss function through slicing. I suppose this workaround confirms that there is something I don’t understand about wrapping them in a `Vector`.
---
### Author:  [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)
#### Post date: [June 23, 2021, 8:52pm UTC](https://discourse.julialang.org/t/flux-jl-inconsistent-training-on-custom-architecture/63326/10 "2021-06-23T20:52:12Z")
I was referring to the MWE, in which only `model.K` is used. Incidentally, the gradient there is `nothing` as well in the MWE because of an interaction between implicit parameters (i.e. the thing you get from `params` and pass to `gradient`) and the AD (Zygote). c.f:
```julia-repl
julia> using Zygote, LinearAlgebra
julia> X = [rand(2, 2)]
1-element Vector{Matrix{Float64}}:
[0.6202008269430048 0.8555679159356662; 0.8423362289177463 0.09771425479926421]
# 1. this doesn't work
julia> gradient(() -> norm(X[1]), Params(X)).grads
IdDict{Any, Any} with 2 entries:
[0.620201 0.855568; 0.842336 0.0977143] => nothing # gradient wrt. X[1]
:(Main.X) => Union{Nothing, Matrix{Float64}}[[0.45775 0.631467; 0.621701 0.0721198]]
# 2. but this does
julia> gradient(() -> norm(X[1]), Params([X])).grads
IdDict{Any, Any} with 2 entries:
:(Main.X) => Union{Nothing, Matrix{Float64}}[[0.45775 0.631467; 0.621701 0.0721198]]
[[0.620201 0.855568; 0.842336 0.0977143]] => Union{Nothing, Matrix{Float64}}[[0.45775 0.631467; 0.621701 0.0721198]]
# 3. as does this
julia> x₁ = X[1]
2×2 Matrix{Float64}:
0.620201 0.855568
0.842336 0.0977143
julia> gradient(() -> norm(x₁), Params(X)).grads
IdDict{Any, Any} with 2 entries:
[0.620201 0.855568; 0.842336 0.0977143] => [0.45775 0.631467; 0.621701 0.0721198]
:(Main.x₁) => [0.45775 0.631467; 0.621701 0.0721198]
# 4. and this (note explicit instead of implicit parameters.
# That is, we pass X directly and use it instead of params(X)).
# This works with full Flux models too!
julia> gradient(x -> norm(x[1]), X)[1]
1-element Vector{Union{Nothing, Matrix{Float64}}}:
[0.4577503206426886 0.6314672132598408; 0.6217013298363201 0.07211975463850043]
```
The gist is that `params` (and `Params`, if given a single argument) splat their arguments into an underlying `IdDict`:
```julia-repl
julia> params(X).order[1]
2×2 Matrix{Float64}: # this is X[1], where you'd expect it to be X itself
0.620201 0.855568
0.842336 0.0977143
```
For whatever reason, Zygote isn’t smart enough to link the `X[1]` in the loss to the actual value of `X[1]` in the params. You can see I avoid this in 2) and 3) by stopping `X` from being unravelled and hoisting the declaration of `X[1]` into a variable respectively.