How do you deserialize stateful optimizers in Flux?

Specific Domains Machine Learning
1

For some stateful optimizers, like ADAM, it has an IdDict field to store momentums of parameters. When deserializing the optimizer with BSON, the id info is lost in the newly created IdDict. This makes it hard to truly recover from some checkpoints.

(I think in TF, each variable has a name so things may be easier?)

2

This should be fixed in latest BSON release

1 Like
3

Ah, thanks for reminding!

I just find the trick here is to save the model and the optimizer into a single bson file.

4

Why? Won’t saving them separately work the same?

5

Yeah, I’m afraid not.

julia> using BSON

julia> a = [1]
1-element Array{Int64,1}:
 1

julia> b = IdDict()
IdDict{Any,Any}b with 0 entries

julia> b[a] = 2
2

julia> b
IdDict{Any,Any} with 1 entry:
  [1] => 2

julia> BSON.@save "b.bson" b

julia> BSON.@save "a.bson" a
 
julia> BSON.@load "a.bson" a

julia> BSON.@load "b.bson" b

julia> a
1-element Array{Int64,1}:
 1

julia> b
IdDict{Any,Any} with 1 entry:
  [1] => 2

julia> b[a]
ERROR: KeyError: key [1] not found
...
2 Likes