Conversation
Summary of ChangesHello @milankl, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces the foundational elements for loading PyTorch models, specifically those saved in the NPZ format, into the Lux deep learning framework. It establishes the necessary package dependencies and a new extension that provides a basic mechanism to read NPZ files, reconstruct a predefined Lux model architecture, and populate its weights. This is an initial draft aimed at exploring the integration of external model formats. Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces initial support for loading PyTorch models from NPZ files into Lux.jl. It correctly updates the Project.toml to include the NPZ dependency and adds a new extension module NPZExt. The NPZExt.jl file defines a DataLoader abstraction and an NPZLoader to handle the loading process. However, there are critical issues related to undefined variables and significant limitations due to hardcoded model architectures, which need to be addressed for this feature to be generally usable.
| lux_layer_params.weight .= Float32.(weights[py_name * ".weight"]) | ||
| lux_layer_params.bias .= Float32.(weights[py_name * ".bias"]) |
There was a problem hiding this comment.
The variable weights is used to access the loaded parameters, but it is not defined within the Lux.load function. The result of NPZ.npzread(filepath) is assigned to the file variable. This will lead to an UndefVarError when trying to access weights[py_name * ".weight"] and weights[py_name * ".bias"].
The correct variable to use here should be file.
lux_layer_params.weight .= Float32.(file[py_name * ".weight"])
lux_layer_params.bias .= Float32.(file[py_name * ".bias"])| # TODO don't hardcode | ||
| layer_map = [ | ||
| "embed_layer" => :layer_1, | ||
| "layer_1" => :layer_2, | ||
| "layer_2" => :layer_4, | ||
| "layer_3" => :layer_6, | ||
| "output_layer" => :layer_7 | ||
| ] |
There was a problem hiding this comment.
The layer_map is currently hardcoded for a very specific model architecture. This significantly limits the utility of the NPZLoader as it can only load models that precisely match this predefined mapping. For a general PyTorch model loading feature, this mapping should ideally be dynamic, inferred from the NPZ file metadata, or provided as a configurable input.
This hardcoding makes the loader non-reusable for other PyTorch models.
| # TODO don't hardcode | ||
| 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) | ||
| ) |
There was a problem hiding this comment.
Similar to the layer_map, the model architecture itself is hardcoded. This means the NPZLoader can only load weights into this exact Chain structure. To support loading diverse PyTorch models, the Lux model architecture should be constructed dynamically based on information extracted from the NPZ file or provided by the user, rather than being fixed.
This hardcoding prevents the loader from being a general-purpose utility.
| @kwdef struct NPZLoader <: PytorchLoader | ||
| dummy_field::Bool = true | ||
| end |
There was a problem hiding this comment.
The NPZLoader struct includes a dummy_field that is always true and doesn't appear to serve any configuration purpose. If there are no specific fields required for NPZLoader configuration, it can be defined as a simple struct without fields to improve clarity and avoid unnecessary data.
Consider defining it as struct NPZLoader <: PytorchLoader end if no fields are needed.
struct NPZLoader <: PytorchLoader end| end | ||
|
|
||
| # fallback | ||
| Lux.load(::DataLoader, args...) = error("DataLoader $(typeof(args[1])) not implemented.") |
There was a problem hiding this comment.
The fallback error message error("DataLoader $(typeof(args[1])) not implemented.") might be less informative than intended. If args contains multiple arguments, args[1] might not always be the DataLoader instance itself. It would be clearer to explicitly refer to the type of the DataLoader argument N.
For example, error("DataLoader $(typeof(N)) not implemented.") would be more precise.
Lux.load(N::DataLoader, args...) = error("DataLoader $(typeof(N)) not implemented.")1 participant
Just an initial draft following #1657. Current questions
using NPZfor npz okay but what about .pt?