# How to read model weights using FluxTraining (stateaccess issues)?

**URL:** https://discourse.julialang.org/t/how-to-read-model-weights-using-fluxtraining-stateaccess-issues/106409
**Category:** Machine Learning
**Tags:** question, flux
**Created:** [November 18, 2023, 3:44pm UTC](https://discourse.julialang.org/t/how-to-read-model-weights-using-fluxtraining-stateaccess-issues/106409 "2023-11-18T15:44:56Z")
**Posts on this page:** 4
**Page:** 1

<div class="post-metadata">

### Author: ![neuromancer](https://avatars.discourse-cdn.com/v4/letter/n/46a35a/32.png) [@neuromancer](https://discourse.julialang.org/u/neuromancer)
#### Post date: [November 18, 2023, 3:44pm UTC](https://discourse.julialang.org/t/how-to-read-model-weights-using-fluxtraining-stateaccess-issues/106409/1 "2023-11-18T15:44:56Z")

</div>

## Summary

I am looking for some help with implementing a custom validation loop using `FluxTraining`. Briefly, I am using `Flux` to optimize a matrix to satisfy a data-driven cost-function. Since the cost function depends on the data, it’s better for me to use an ML framework than a regular optimization package. The model is a single matrix, so in Flux parlance:

```julia-auto
model = Dense(m => n, identity; bias=false)

```

Due to the structure of my loss function and evaluation metrics, I need to access the weights (the matrix) during computation of the loss and evaluation.

I am running into an error where I am being blocked _by FluxTraining_ from accessing the mode weights during eval (but not training).

```julia-auto
Epoch 1 MyValidationPhase() ...
ERROR: LoadError: FluxTraining.ProtectedException("Read access to Learner.model disallowed.")
Stacktrace:
  [1] getfieldperm(data::Learner, field::Symbol, perm::Nothing)
    @ FluxTraining ~/.julia/packages/FluxTraining/xCOPx/src/callbacks/protect.jl:63
  [2] getproperty(protected::FluxTraining.Protected{Learner}, field::Symbol)
    @ FluxTraining ~/.julia/packages/FluxTraining/xCOPx/src/callbacks/protect.jl:18
  [3] step!(metric::ReconstructionMetric{Number}, learner::FluxTraining.Protected{Learner}, phase::MyValidationPhase)
    @ Main ~/.guild/runs/1d889bb02af249a1a4fcf77833cc2296/train.jl:260
  [4] on(#unused#::FluxTraining.Events.StepEnd, phase::MyValidationPhase, metrics::Metrics, learner::FluxTraining.Protected{Learner})
    @ FluxTraining ~/.julia/packages/FluxTraining/xCOPx/src/callbacks/metrics.jl:74
  [5] _on(e::FluxTraining.Events.StepEnd, p::MyValidationPhase, cb::Metrics, learner::Learner)
    @ FluxTraining ~/.julia/packages/FluxTraining/xCOPx/src/callbacks/callback.jl:254
  [6] handle(runner::FluxTraining.LinearRunner, event::FluxTraining.Events.StepEnd, phase::MyValidationPhase, learner::Learner)
    @ FluxTraining ~/.julia/packages/FluxTraining/xCOPx/src/callbacks/execution.jl:12
  [7] (::FluxTraining.var"#handlefn#82"{Learner, MyValidationPhase})(e::FluxTraining.Events.StepEnd)
    @ FluxTraining ~/.julia/packages/FluxTraining/xCOPx/src/training.jl:129
  [8] runstep(stepfn::var"#9#10"{Learner}, learner::Learner, phase::MyValidationPhase, initialstate::NamedTuple{(:xs, :ys), Tuple{Matrix{Float32}, Matrix{Float32}}})
    @ FluxTraining ~/.julia/packages/FluxTraining/xCOPx/src/training.jl:134
  [9] step!(learner::Learner, phase::MyValidationPhase, batch::Tuple{Matrix{Float32}, Matrix{Float32}})
    @ Main ~/.guild/runs/1d889bb02af249a1a4fcf77833cc2296/train.jl:301
 [10] (::FluxTraining.var"#71#72"{Learner, MyValidationPhase, MLUtils.DataLoader{Tuple{Matrix{Float32}, Matrix{Float32}}, Random._GLOBAL_RNG, Val{nothing}}})(#unused#::Function)
    @ FluxTraining ~/.julia/packages/FluxTraining/xCOPx/src/training.jl:24
 [11] runepoch(epochfn::FluxTraining.var"#71#72"{Learner, MyValidationPhase, MLUtils.DataLoader{Tuple{Matrix{Float32}, Matrix{Float32}}, Random._GLOBAL_RNG, Val{nothing}}}, learner::Learner, phase::MyValidationPhase)
    @ FluxTraining ~/.julia/packages/FluxTraining/xCOPx/src/training.jl:105
 [12] epoch!(learner::Learner, phase::MyValidationPhase, dataiter::MLUtils.DataLoader{Tuple{Matrix{Float32}, Matrix{Float32}}, Random._GLOBAL_RNG, Val{nothing}})
    @ FluxTraining ~/.julia/packages/FluxTraining/xCOPx/src/training.jl:22
 [13] main()
    @ Main ~/.guild/runs/1d889bb02af249a1a4fcf77833cc2296/train.jl:449
 [14] top-level scope
    @ ~/.guild/runs/1d889bb02af249a1a4fcf77833cc2296/train.jl:453
in expression starting at /home/alec/.guild/runs/1d889bb02af249a1a4fcf77833cc2296/train.jl:453

```

## System/Julia information

I am running on Pop!\_OS 22.04 with Julia 1.9.0. All code is being run on the CPU.  
[Here is my Manifest.toml file](https://gist.github.com/alec-hoyland/48002cc51aa8c1bab87a76b639ffa0be).

## The Training Phase (this works)

My custom training phase is working fine.

```julia-auto

function pack_loss(state::FluxTraining.PropDict, loss_names, loss_vals)
    state.loss = reduce(+, loss_vals)
    for i in eachindex(loss_names)
        setproperty!(state, Symbol(loss_names[i]), loss_vals[i])
    end
    return state
end

function unpack_loss(state::FluxTraining.PropDict, loss_name)
    return getproperty(state, Symbol(loss_name))
end

"""
Custom training phase
that exposes the model parameters.
"""
struct MyTrainingPhase <: FluxTraining.AbstractTrainingPhase end

"""
    FluxTraining.step!(learner, phase::MyTrainingPhase, batch)

`step!` function for the custom training phase.
The `runstep` function returns the `state`, but the `state` is also packaged
into `learner.step`, which OnlineStats.means that we can access it later.

This function handles the custom loss function with keyword arguments
and also handles the state, adding the loss components to the `state`.
"""
function FluxTraining.step!(learner, phase::MyTrainingPhase, batch)
    xs, ys = batch

    # `runstep` has the signature: runstep(stepfn, learner, phase, batch)
    # where `stepfn` has the signature: stepfn(handle, state::PropDict)
    FluxTraining.runstep(learner, phase, (; xs = xs, ys = ys)) do handle, state
        # `stepfn` calls the `_gradient` method which is a fallback
        # for different formats of Flux models,
        # but requires a lossfn as the first input, which has the signature: lossfn(model)
        state.grads = FluxTraining._gradient(learner.optimizer,
            learner.model,
            learner.params) do model
            # get the model outputs
            state.ŷs = model(state.xs)
            handle(FluxTraining.LossBegin())
            # compute the loss using the custom loss function
            loss_names, loss_vals = learner.lossfn(state.ŷs, state.ys; m = learner.model)
            # pack the loss in the state
            state = pack_loss(state, loss_names, loss_vals)
            handle(FluxTraining.BackwardBegin())
            # return the total loss
            return state.loss
        end
        handle(BackwardEnd())
        # update parameters
        learner.params, learner.model = FluxTraining._update!(learner.optimizer,
            learner.params,
            learner.model, state.grads)
    end
end

```

The main differences between this and the standard training phase baked into `FluxTraining` are that the loss function returns multiple components of the loss, and that the loss function has the signature `lossfn(y_hat, y; m)` where `m` is the model.

## The Validation Phase (this is broken)

> For reasons of protecting unpublished work, my advisor has asked me to obscure some of the details here, but I am certain that nothing I am leaving out is causing the problem. Thank you for understanding.

Here’s the validation phase code. I’m pretty sure this is working fine.

```julia-auto
struct MyValidationPhase <: FluxTraining.AbstractValidationPhase end

function FluxTraining.step!(learner, phase::MyValidationPhase, batch)
    xs, ys = batch
    FluxTraining.runstep(learner, phase, (; xs = xs, ys = ys)) do _, state
        state.ŷs = learner.model(state.xs)
        loss_names, loss_vals = learner.lossfn(state.ŷs, state.ys; m = learner.model)
        state = pack_loss(state, loss_names, loss_vals)
    end
end

```

Here’s the code for the custom eval metrix callback that’s causing me problems…

```julia-auto
mutable struct ReconstructionMetric{T} <: FluxTraining.AbstractMetric
    decisionfn::Any
    decisionfn_kwargs::NamedTuple
    reconstructionfn::Any
    reconstructionfn_kwargs::NamedTuple
    statistic::OnlineStats.OnlineStat{T}
    _statistic::Any
    name::Any
    device::Any
    P::Any
    last::Union{Nothing, T}
    target_signal_repr::Any
end

function FluxTraining.reset!(metric::ReconstructionMetric{T}) where T
    metric.statistic = deepcopy(metric._statistic)
end

function FluxTraining.step!(metric::ReconstructionMetric, learner, phase)
    if phase isa metric.P
        y = metric.decisionfn(learner.model.weight,
            metric.target_signal_repr;
            metric.decisionfn_kwargs...)
        x = metric.reconstructionfn(y,
            learner.model.weight;
            metric.reconstructionfn_kwargs...)
        metric.last = StatsBase.cor(x, metric.target_signal_repr)
        OnlineStats.fit!(metric.statistic, metric.last)
    else
        metric.last = nothing
    end
end

# Here is where I think the problem is.
# I am getting blocked from accessing learner.model.weight,
# even though I am in theory using the `stateaccess` named tuple correctly...afaict.
function FluxTraining.stateaccess(::ReconstructionMetric)
    (
        model = FluxTraining.Read(),
        params = FluxTraining.Read(),
        cbstate = (metricsstep = FluxTraining.Write(), metricsepoch = FluxTraining.Write(), history = FluxTraining.Read()),
        step = FluxTraining.Read()
    )
end

FluxTraining.runafter(::ReconstructionMetric) = (FluxTraining.Recorder,)
FluxTraining.stepvalue(metric::ReconstructionMetric) = metric.last
FluxTraining.metricname(metric::ReconstructionMetric) = metric.name

function FluxTraining.epochvalue(metric::ReconstructionMetric)
    if isnothing(metric.last)
        nothing
    else
        OnlineStats.value(metric.statistic)
    end
end

```

Any advice would be appreciated!

---

<div class="post-metadata">

### Author: ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)
#### Post date: [November 20, 2023, 5:59pm UTC](https://discourse.julialang.org/t/how-to-read-model-weights-using-fluxtraining-stateaccess-issues/106409/2 "2023-11-20T17:59:51Z")

</div>

It looks like you have enough to create a MWE without revealing any unpublished details. Doing that would help immensely with tracking down what the issue is. As-is, having to mentally trace through the execution of each code path with just the snippets you’ve provided is a bit too much effort.

---

<div class="post-metadata">

### Author: ![neuromancer](https://avatars.discourse-cdn.com/v4/letter/n/46a35a/32.png) [@neuromancer](https://discourse.julialang.org/u/neuromancer)
#### Post date: [November 21, 2023, 1:58pm UTC](https://discourse.julialang.org/t/how-to-read-model-weights-using-fluxtraining-stateaccess-issues/106409/3 "2023-11-21T13:58:46Z")

</div>

Fair enough, here is a MWE. Any help you can provide would be appreciated. Thanks!

```julia
using Flux
using FluxTraining
using OnlineStats
using LinearAlgebra

"""
Custom type for my metric.
This is basically a duplicate of the standard Metric.
"""
mutable struct MyMetric{T} <: FluxTraining.AbstractMetric
    statistic::OnlineStats.OnlineStat{T}
    _statistic::Any
    name::Any
    device::Any
    P::Any
    last::Union{Nothing, T}
end

"""
Outer constructor for MyMetric.
"""
function MyMetric(;
        statistic = OnlineStats.Mean(Float32),
        device = cpu,
        phase = ValidationPhase,
        name = "MyMetric")
    return MyMetric(statistic, deepcopy(statistic), name, device, phase, nothing)
end

"""
Reset MyMetric back to the initial value.
"""
function FluxTraining.reset!(metric::MyMetric{T}) where T
    metric.statistic = deepcopy(metric._statistic)
end

"""
We will use the L1 norm as an "example function"
that requires access to model weights.
"""
function l1_metric(W::Matrix)
    return norm(W, 1) / size(W, 1)
end

"""
Compute the metric by taking the L1 norm of the model weight matrix.
"""
function FluxTraining.step!(metric::MyMetric, learner, phase)
    if phase isa metric.P
        metric.last = l1_metric(learner.model.weight)
        OnlineStats.fit!(metric.statistic, metric.last)
    else
        metric.last = nothing
    end
end

function Base.show(io::IO, metric::MyMetric{T}) where {T}
    print(io, "Metric(", metric.name, ")")
end

FluxTraining.runafter(::MyMetric) = (FluxTraining.Recorder,)
FluxTraining.stepvalue(metric::MyMetric) = metric.last
FluxTraining.metricname(metric::MyMetric) = metric.name

function FluxTraining.epochvalue(metric::MyMetric)
    if isnothing(metric.last)
        nothing
    else
        OnlineStats.value(metric.statistic)
    end
end

function FluxTraining.stateaccess(::MyMetric)
    return (
        model = FluxTraining.Read(),
        params = FluxTraining.Read(),
        cbstate = (metricsstep = FluxTraining.Write(), metricsepoch = FluxTraining.Write(), history = FluxTraining.Read()),
        step = FluxTraining.Read(),
    )
end

function main()
    in_dim = 10
    out_dim = 1
    n_samples = 64
    model = Dense(in_dim => out_dim, identity; bias=false)

    X = rand(in_dim, n_samples) |> f32
    y = rand(out_dim, n_samples) |> f32
    train_dataloader = Flux.DataLoader((X, y))
    val_dataloader = deepcopy(train_dataloader)

    callbacks = [FluxTraining.Metrics(MyMetric())]
    opt_state = Flux.Adam(1f-4)

    learner = FluxTraining.Learner(model, Flux.mse; callbacks = callbacks, optimizer = opt_state)

    for i = 1:3
        FluxTraining.epoch!(learner, FluxTraining.TrainingPhase(), train_dataloader)
        FluxTraining.epoch!(learner, ValidationPhase(), val_dataloader)
    end
end

```

---

<div class="post-metadata">

### Author: ![neuromancer](https://avatars.discourse-cdn.com/v4/letter/n/46a35a/32.png) [@neuromancer](https://discourse.julialang.org/u/neuromancer)
#### Post date: [January 9, 2024, 9:35pm UTC](https://discourse.julialang.org/t/how-to-read-model-weights-using-fluxtraining-stateaccess-issues/106409/4 "2024-01-09T21:35:21Z")

</div>

The issue is that `stateaccess` need to be for `::Metrics` not `::MyMetric`.

```julia
function FluxTraining.stateaccess(::FluxTraining.Metrics)
    (model = FluxTraining.Read(),
        params = FluxTraining.Read(),
        cbstate = (metricsstep = FluxTraining.Write(),
            metricsepoch = FluxTraining.Write(),
            history = FluxTraining.Read()),
        step = FluxTraining.Read())
end

```