huggingface / huggingface/candle

question about VarBuilder:from_pth

Open
#2,990 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Rust
Stars
21k
Forks
1.8k
Avg merge
16h 42m
Merged PRs (30d)
25

Description

Ні,
I'm trying to load a .pth file using the following code, but the variable vb remains empty.
```rust
let device = Device:Cpu;
let vb = VarBuilder:from_pth(pth_file, Type::F32, &device)?;
```
The structure of my .pth file is:
`Dict([("model", Dict([("bias", Reduce), ("weight", Reduce), ...]))])`
It seems `VarBuilder:from_pth` currently only supports the format:
`Dict([('bias", Reduce), ("weight", Reduce), ...])`

I made the following changes:
1. Added `Object::Dict(key_values) => key_values[0].1.clone()` in `candle_core::pickle::read_pth_tensor_info`
2. Changed `tensor_infos` and `path` fields in the `PthTensors` struct to pub
These modifications allowed the PTH file to load correctly.
```rust
let obj = match obj {
Object::Build { callable, args } => match *callable {
Object::Reduce { callable, args: _ } => match *callable {
Object::Class {
module_name,
class_name,
} if module_name == "__torch__" && class_name == "Module" => *args,
_ => continue,
},
_ => continue,
},
Object::Dict(key_values) => key_values[0].1.clone(), // I add this line
obj => obj,
};
```
I'm wondering if there's a better/more efficient way to load my pth file?

pth file example:
```json
Dict(
[
(
Unicode(
"model",
),
Dict(
[
(
Unicode(
"ref_enc.layernorm.bias",
),
Reduce {
callable: Class {
module_name: "torch._utils",
class_name: "_rebuild_tensor_v2",
},
args: Tuple(
[
PersistentLoad(
Tuple(
[
Unicode(
"storage",
),
Class {
module_name: "torch",
class_name: "FloatStorage",
},
Unicode(
"482",
),
Unicode(
"cpu",
),
Int(
513,
),
],
),
),
Int(
0,
),
Tuple(
[
Int(
513,
),
],
),
Tuple(
[
Int(
1,
),
],
),
Bool(
false,
),
Dict(
[],
),
],
),
},
),
(
Unicode(
"ref_enc.layernorm.weight",
),
Reduce {
callable: Class {
module_name: "torch._utils",
class_name: "_rebuild_tensor_v2",
},
args: Tuple(
[
PersistentLoad(
Tuple(
[
Unicode(
"storage",
),
Class {
module_name: "torch",
class_name: "FloatStorage",
},
Unicode(
"481",
),
Unicode(
"cpu",
),
Int(
513,
),
],
),
),
Int(
0,
),
Tuple(
[
Int(
513,
),
],
),
Tuple(
[
Int(
1,
),
],
),
Bool(
false,
),
Dict(
[],
),
],
),
},
),
```

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with candle_core::pickle::read_pth_tensor_info and the PthTensors fields mentioned in the issue, then inspect how VarBuilder::from_pth handles the parsed structure. Verify behavior against the nested model dictionary shown in the example; done means the PTH file loads its tensors without requiring public internal fields.

Written by the indexing model from the issue text.

Assessment

Tech stack
pytorch, rust
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.