Struct tch::nn::VarStore

pub struct VarStore {
    pub variables_: Arc<Mutex<Variables>>,
    // some fields omitted
}

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 new(device: Device) -> VarStore

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

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 gets saved in the given file.

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 gets 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_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 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.