Expand description
A VarStore is used to store variables used by one or multiple layers. It specifies a single device where all variables are stored.
Fields
variables_: Arc<Mutex<Variables>>
Implementations
sourceimpl VarStore
impl VarStore
sourcepub fn new(device: Device) -> VarStore
pub fn new(device: Device) -> VarStore
Creates a new var-store located on the specified device.
sourcepub fn is_empty(&self) -> bool
pub fn is_empty(&self) -> bool
Returns true if no tensors are currently stored on this var-store.
sourcepub fn trainable_variables(&self) -> Vec<Tensor>
pub fn trainable_variables(&self) -> Vec<Tensor>
Returns all the trainable variables for this var-store.
sourcepub fn variables(&self) -> HashMap<String, Tensor>
pub fn variables(&self) -> HashMap<String, Tensor>
Returns all variables along with their names.
sourcepub fn root(&self) -> Path<'_>
pub fn root(&self) -> Path<'_>
Gets the root path for this variable store.
Variables are named and organized using paths. This function returns the top level path for the var store and can be combined with ‘/’ to create sub-paths.
sourcepub fn save<T: AsRef<Path>>(&self, path: T) -> Result<(), TchError>
pub fn save<T: AsRef<Path>>(&self, path: T) -> Result<(), TchError>
Saves the var-store variable values to a file.
Weight values for all the tensors currently stored in the var-store are saved in the given file.
sourcepub fn save_to_stream<W: Write>(&self, stream: W) -> Result<(), TchError>
pub fn save_to_stream<W: Write>(&self, stream: W) -> Result<(), TchError>
Saves the var-store variable values to a stream.
Weight values for all the tensors currently stored in the var-store gets saved in the given stream.
sourcepub fn load<T: AsRef<Path>>(&mut self, path: T) -> Result<(), TchError>
pub fn load<T: AsRef<Path>>(&mut self, path: T) -> Result<(), TchError>
Loads the var-store variable values from a file.
Weight values for all the tensors currently stored in the var-store are loaded from the given file. Note that the set of variables stored in the var-store is not changed, only the values for these tensors are modified.
sourcepub fn load_from_stream<S: Read + Seek>(
&mut self,
stream: S
) -> Result<(), TchError>
pub fn load_from_stream<S: Read + Seek>(
&mut self,
stream: S
) -> Result<(), TchError>
Loads the var-store variable values from a stream.
Weight values for all the tensors currently stored in the var-store gets loaded from the given stream. Note that the set of variables stored in the var-store is not changed, only the values for these tensors are modified.
sourcepub fn load_partial<T: AsRef<Path>>(
&mut self,
path: T
) -> Result<Vec<String>, TchError>
pub fn load_partial<T: AsRef<Path>>(
&mut self,
path: T
) -> Result<Vec<String>, TchError>
Loads the var-store variable values from a file if it exists.
Weight values for the tensors currently stored in the var-store and the given file get loaded from the given file. If a variable in the var store is not present in the given file, it is skipped and its values are not updated. This method should be used if pre-trained weight for only parts of the model are available. Note that the set of variables stored in the var-store is not changed, only the values for these tensors are modified.
Returns a String Vector containing the names of missing variables.
sourcepub fn freeze(&mut self)
pub fn freeze(&mut self)
Freezes a var store.
Gradients for the variables in this store are not tracked anymore.
sourcepub fn unfreeze(&mut self)
pub fn unfreeze(&mut self)
Unfreezes a var store.
Gradients for the variables in this store are tracked again.
sourcepub fn set_kind(&mut self, kind: Kind)
pub fn set_kind(&mut self, kind: Kind)
Casts all variables in a var store to the target kind .
For floating-point conversion, methods half
, bfloat16
, float
and double
should be preferred as they ensure only float-like variables will be converted
to the target type.
sourcepub fn half(&mut self)
pub fn half(&mut self)
Casts all float-like variable of a var store to half-precision (Half kind).
sourcepub fn bfloat16(&mut self)
pub fn bfloat16(&mut self)
Casts all float-like variable of a var store to bfloat16-precision (BFloat16 kind).
sourcepub fn float(&mut self)
pub fn float(&mut self)
Casts all float-like variable of a var store to single-precision (Float kind).
sourcepub fn double(&mut self)
pub fn double(&mut self)
Casts all float-like variable of a var store to single-precision (Double kind).
sourcepub fn set_device(&mut self, device: Device)
pub fn set_device(&mut self, device: Device)
Migrates a var store and all its tensor to a target device.
Trait Implementations
Auto Trait Implementations
impl RefUnwindSafe for VarStore
impl Send for VarStore
impl Sync for VarStore
impl Unpin for VarStore
impl UnwindSafe for VarStore
Blanket Implementations
sourceimpl<T> BorrowMut<T> for T where
T: ?Sized,
impl<T> BorrowMut<T> for T where
T: ?Sized,
const: unstable · sourcefn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
Mutably borrows from an owned value. Read more