Flux save and load model

New to Julia
1

Hello
I’m saving the UNet model as
function save_unet_train_state(filename::String, model::Unet, opt, loss::Function) @assert endswith(filename, ".jld2") @save filename model opt loss
where opt is opt = RMSProp(learning_rate, momentum)
but on loading, i wrote `

> model_loaded, opt_loaded, loss_loaded = UNet.load_unet(joinpath(pwd(), "UNet.jld2"))
> ┌ Warning: saved type Unet is missing field labels in workspace type; reconstructing
> └ @ JLD2 ~/.julia/packages/JLD2/k9Gt0/src/data/reconstructing_datatypes.jl:152
> ERROR: MethodError: no method matching iterate(::JLD2.ReconstructedTypes.var"##Unet#298")
> Closest candidates are:
>   iterate(::Union{LinRange, StepRangeLen}) at range.jl:664
>   iterate(::Union{LinRange, StepRangeLen}, ::Int64) at range.jl:664

Not sure how to fix this?
Also, the field label is a vector of classes, e.g. [1,2,3]. I thought saving model will save it as model.labels automatically?

2

You might have to run all of the using and includes that were present when you saved the module, so that the deserialiser knows how to construct the objects.

EDIT: I am not familiar with UNet, but shouldn’t you be using a @load to load the data back? Similar to how it is saved in the first place?

3

Sorry I didn’t include the original definition of load.
@load filename model opt loss labels
However, the @save didn’t save one field which is labels.
Labels was added at training when model was initialized
model = Unet(num_channels, labels) |> device