huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-19 19:58:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

VarBuilder cleanup (#627)

Browse Source 
* VarBuilder cleanup.

* Implement the basic varbuilders.

* Add the sharded code.

* Proper support for tensor sharding.
This commit is contained in:
Laurent Mazare
2023-08-27 18:03:26 +01:00
committed by GitHub
parent be471d50ab
commit 4c338b0cd9
12 changed files with 409 additions and 291 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
candle-nn/src/conv.rs
View File

@ -124,7 +124,7 @@ pub fn conv1d(
vs: crate::VarBuilder,
) -> Result<Conv1d> {
let init_ws = crate::init::DEFAULT_KAIMING_NORMAL;
let ws = vs.get_or_init(
let ws = vs.get_with_hints(
(out_channels, in_channels / cfg.groups, kernel_size),
"weight",
init_ws,
@ -134,7 +134,7 @@ pub fn conv1d(
lo: -bound,
up: bound,
};
let bs = vs.get_or_init(out_channels, "bias", init_bs)?;
let bs = vs.get_with_hints(out_channels, "bias", init_bs)?;
Ok(Conv1d::new(ws, Some(bs), cfg))
}
@ -146,7 +146,7 @@ pub fn conv2d(
vs: crate::VarBuilder,
) -> Result<Conv2d> {
let init_ws = crate::init::DEFAULT_KAIMING_NORMAL;
let ws = vs.get_or_init(
let ws = vs.get_with_hints(
(
out_channels,
in_channels / cfg.groups,
@ -161,7 +161,7 @@ pub fn conv2d(
lo: -bound,
up: bound,
};
let bs = vs.get_or_init(out_channels, "bias", init_bs)?;
let bs = vs.get_with_hints(out_channels, "bias", init_bs)?;
Ok(Conv2d::new(ws, Some(bs), cfg))
}
@ -173,7 +173,7 @@ pub fn conv2d_no_bias(
vs: crate::VarBuilder,
) -> Result<Conv2d> {
let init_ws = crate::init::DEFAULT_KAIMING_NORMAL;
let ws = vs.get_or_init(
let ws = vs.get_with_hints(
(
out_channels,
in_channels / cfg.groups,