Skip to main content

Model

Trait Model 

Source
pub trait Model: Display {
    type Input;
    type Output;

    // Required methods
    fn forward(
        &mut self,
        input: Self::Input,
    ) -> Result<Self::Output, ModelError>;
    fn require_grad(&mut self, grad_enabled: bool) -> Result<(), ModelError>;
    fn to_device(&mut self, device: &Device) -> Result<(), ModelError>;
    fn to_dtype(&mut self, dtype: &DType) -> Result<(), ModelError>;
    fn parameters(&self) -> Result<Vec<Tensor>, ModelError>;
    fn named_parameters(&self) -> Result<HashMap<String, Tensor>, ModelError>;

    // Provided methods
    fn save(&self, file_path: &str) -> Result<(), ModelError> { ... }
    fn load(
        &mut self,
        file_path: &str,
        device: &Device,
    ) -> Result<(), ModelError> { ... }
}

Required Associated Types§

Required Methods§

Source

fn forward(&mut self, input: Self::Input) -> Result<Self::Output, ModelError>

Run the model forward pass.

§Arguments
  • input - The input tensor.
§Returns
  • Ok(Self::Output) - The output tensor if successful.
  • Err(ModelError) - The error when running the model forward pass.
Source

fn require_grad(&mut self, grad_enabled: bool) -> Result<(), ModelError>

Set whether to enable gradient tracking for the model.

§Arguments
  • grad_enabled - Whether to enable gradient tracking for the model.
§Returns
  • Ok(()) - If successful.
  • Err(ModelError) - The error when setting the gradient tracking.
Source

fn to_device(&mut self, device: &Device) -> Result<(), ModelError>

Move the model to the specified device.

§Arguments
  • device - The device to move the model to.
§Returns
  • Ok(()) - If successful.
  • Err(ModelError) - The error when moving the model to the device.
Source

fn to_dtype(&mut self, dtype: &DType) -> Result<(), ModelError>

Convert the model to the specified data type.

§Arguments
  • dtype - The data type to convert the model to.
§Returns
  • Ok(()) - If successful.
  • Err(ModelError) - The error when converting the model to the data type.
Source

fn parameters(&self) -> Result<Vec<Tensor>, ModelError>

Get the parameters of the model.

§Returns
  • Ok(Vec<Tensor>) - The parameters of the model.
  • Err(ModelError) - The error when getting the parameters of the model.
Source

fn named_parameters(&self) -> Result<HashMap<String, Tensor>, ModelError>

Get the named parameters of the model.

§Returns
  • Ok(HashMap<String, Tensor>) - The named parameters of the model.
  • Err(ModelError) - The error when getting the named parameters of the model.

Provided Methods§

Source

fn save(&self, file_path: &str) -> Result<(), ModelError>

Save the model parameters to a file.

§Arguments
  • file_path - The path to the file.
§Returns
  • Ok(()) - If successful.
  • Err(ModelError) - The error when saving the model parameters to the file.
Source

fn load(&mut self, file_path: &str, device: &Device) -> Result<(), ModelError>

Load the model parameters from a file.

§Arguments
  • file_path - The path to the file.
  • device - The device to load the model parameters to.
§Returns
  • Ok(()) - If successful.
  • Err(ModelError) - The error when loading the model parameters from the file.

Dyn Compatibility§

This trait is dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors§