Skip to content

Conversation

@CarloLucibello
Copy link
Member

@CarloLucibello CarloLucibello commented Apr 25, 2023

Fix performance issue mentioned in #2239 (comment)

When loading from a flux model into another everything is fine:

using Metalhead, Flux

m1 = ResNet(18)
m2 = ResNet(18)
@time Flux.loadmodel!(m2, m1) # warmup
@time Flux.loadmodel!(m2, m1)
#  0.003388 seconds (23.01 k allocations: 2.157 MiB) # fast on both master and this PR

When loading from a similar struct such as nested named tuples instead (as proposed in #2239) perf on master is very bad:

using Functors

function state(x)
    if Functors.isleaf(x)
        return x
    else
        return map(state, Functors.children(x))
    end
end

s = state(m1)
@time Flux.loadmodel!(m2, s) # warmup
@time Flux.loadmodel!(m2, s)
# 69.507712 seconds (1.06 G allocations: 40.744 GiB, 4.74% gc time) # master
# 0.002409 seconds (6.85 k allocations: 1.109 MiB)                               # this PR

The perf problem was due to huge strings being created at each recursion step as part of error messages.

@theabhirath

Also removed the filter argument, which was not documented nor used in practice.

@theabhirath
Copy link
Member

filter is important for Metalhead. The motivation is described in #2041

@CarloLucibello
Copy link
Member Author

I'll revert the filter removal. Can we make it documented API?

@darsnack
Copy link
Member

I don't have an issue with that, but I recall that last time there was some bikeshedding and we opted to keep it private. If it's still an issue, no need to let it hold up this PR.

@CarloLucibello CarloLucibello merged commit 7e329f3 into master Apr 25, 2023
rgobbel pushed a commit to rgobbel/Flux.jl that referenced this pull request Apr 25, 2023
* fix perf issue loadmodel!

* reinstate filter

* cleanup

* cleanup
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants