Load PyTorch model weights into Lux

[gregrmunday]

Hi! As part of some work on SpeedWeather.jl#959 I'm doing with @milankl, I've needed to load model weights from a PyTorch-trained model into Lux (just a simple feed-forward network).

After I save the weights from PyTorch into a .npz, I can load them into Lux using

using Lux, Random, NPZ

rng = Random.default_rng()
dense_model = Chain(
        Dense(13 => 32, leakyrelu),
        Dense(32 => 64, leakyrelu),
        Dropout(0.2),
        Dense(64 => 64, leakyrelu),
        Dropout(0.1),
        Dense(64 => 32, leakyrelu),
        Dense(32 => 1)
    )

parameters, states = Lux.setup(rng, dense_model)
weights = npzread("model_weights.npz")

layer_map = [
    "embed_layer"  => :layer_1,
    "layer_1"      => :layer_2,
    "layer_2"      => :layer_4,
    "layer_3"      => :layer_6,
    "output_layer" => :layer_7
]
for (py_name, lux_sym) in layer_map
    lux_layer_params = getproperty(parameters, lux_sym)
    
    lux_layer_params.weight .= Float32.(weights[py_name * ".weight"])
    lux_layer_params.bias   .= Float32.(weights[py_name * ".bias"])
end

and from there use the model like

test_states = Lux.testmode(states)
model = (u, p, s) -> first(Lux.apply(dense_model, u, p, s))
y = model(zeros(Float32, 13), parameters, test_states)

I'm basically wondering if there's a better way to do the above, and if not would it be worth writing a more simple load function into Lux? It feels like this could be quite a common requirement, especially for people coming from PyTorch and looking for an easy bridge between the two!

[avik-pal]

Currently, manually is pretty much the only way, since the layer storages (/ names of parameters) wont match between Lux and Pytorch. That said, I want to wait for EnzymeAD/Reactant.jl#2065 which should let us import models without any boilerplate

[milankl]

I see this would be import models in the compiled form? But that wouldn't allow me to load in some pretrained weights and make some changes on the Lux level to it, e.g. prune some weights, or compile it to another architecture? I like to think of functionality that allows one to just load in the data and then retain the high level Lux access to it.

What Greg is suggesting of course wouldn't easily scale to more complex NN architectures but one could start porting MLPs between torch and Lux at least?

[avik-pal]

Ah, I see. Yes, something like that would be a useful addition. We can limit it to work on simple cases like nn.Linear, nn.Conv, nn.Sequential and explicitly error if we encounter types we don't know about

[milankl]

Yes that's what I thought. Start simple and extend this functionality further if needed. Like

model, parameters, states = Lux.load("model_weights.npz"; kwargs...)

with optional directives how to interpret the torch data passed on by kwargs. And then error/warn if the file contents can't be interpreted correctly. Could be in an LuxNPZExt to avoid dependencies or in lib/LuxLoader a subdir package?

[avik-pal]

We should have a first argument as ::PytorchLoader.

Could be in an LuxNPZExt to avoid dependencies or in lib/LuxLoader a subdir package?

Let's do an extension

[gregrmunday]

Thanks both, sounds good to me. I'll crack on with something over the next little bit!