huggingface / huggingface/candle
question about VarBuilder:from_pth
- 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