Struct tch::nn::VarStore

pub struct VarStore {
    pub variables_: Arc<Mutex<Variables>>,
    /* private fields */
}

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

impl VarStore

pub fn read_safetensors<T: AsRef<Path>>(&self, path: T) -> Result<(), TchError>

Read data from safe tensor file, missing tensors will raise a error.

pub fn fill_safetensors<P: AsRef<Path>>(&self, path: P) -> Result<(), TchError>

impl VarStore

pub fn new(device: Device) -> VarStore

Creates a new var-store located on the specified device.

pub fn merge( var_stores: Vec<(VarStore, Option<&str>)> ) -> Result<VarStore, TchError>

pub fn device(&self) -> Device

Gets the device for this var-store.

pub fn len(&self) -> usize

Returns the number of tensors currently stored on this var-store.

pub fn is_empty(&self) -> bool

Returns true if no tensors are currently stored on this var-store.

pub fn trainable_variables(&self) -> Vec<Tensor>

Returns all the trainable variables for this var-store.

pub fn variables(&self) -> HashMap<String, Tensor>

Returns all variables along with their names.

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.

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.

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.

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.

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.

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.

pub fn freeze(&mut self)

Freezes a var store.

Gradients for the variables in this store are not tracked anymore.

pub fn unfreeze(&mut self)

Unfreezes a var store.

Gradients for the variables in this store are tracked again.

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.

pub fn half(&mut self)

Casts all float-like variable of a var store to half-precision (Half kind).

pub fn bfloat16(&mut self)

Casts all float-like variable of a var store to bfloat16-precision (BFloat16 kind).

pub fn float(&mut self)

Casts all float-like variable of a var store to single-precision (Float kind).

pub fn double(&mut self)

Casts all float-like variable of a var store to single-precision (Double kind).

pub fn set_device(&mut self, device: Device)

Migrates a var store and all its tensor to a target device.

pub fn copy(&mut self, src: &VarStore) -> Result<(), TchError>

Copies variable values from a source var store to this var store.

All the variables in this var store have to exist with the same name in the source var store, otherwise an error is returned.

Trait Implementations

impl Debug for VarStore

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.