Struct burn_import::pytorch::LoadArgs

pub struct LoadArgs {
    pub file: PathBuf,
    pub key_remap: Vec<(Regex, String)>,
    pub top_level_key: Option<String>,
    pub debug: bool,
}

Arguments for loading a PyTorch file.

Fields

• file - The path to the file to load.
• key_remap - A vector of tuples containing a regular expression and a replacement string. See regex::Regex::replace for more information.

Notes

Use Netron to inspect the keys of the PyTorch file (.pt extension).

Examples

use burn_import::pytorch::{LoadArgs, PyTorchFileRecorder};
use burn::record::FullPrecisionSettings;
use burn::record::Recorder;

let args = LoadArgs::new("tests/key_remap/key_remap.pt".into())
   .with_key_remap("conv\\.(.*)", "$1"); // // Remove "conv" prefix, e.g. "conv.conv1" -> "conv1"

let record = PyTorchFileRecorder::<FullPrecisionSettings>::default()
  .load(args)
  .expect("Should decode state successfully");

Fields

file: PathBuf

The path to the file to load.

key_remap: Vec<(Regex, String)>

A list of key remappings.

top_level_key: Option<String>

Top-level key to load state_dict from the file. Sometimes the state_dict is nested under a top-level key in a dict.

debug: bool

Whether to print debug information.

Implementations

impl LoadArgs

pub fn new(file: PathBuf) -> Self

Creates a new LoadArgs instance.

Arguments

• file - The path to the file to load.

pub fn with_key_remap(self, pattern: &str, replacement: &str) -> Self

Sets key remapping.

Arguments

• pattern - The Regex pattern to be replaced.
• replacement - The pattern to replace with.

See Regex for the pattern syntax and Replacement for the replacement syntax.

pub fn with_top_level_key(self, key: &str) -> Self

Sets the top-level key to load state_dict from the file. Sometimes the state_dict is nested under a top-level key in a dict.

Arguments

• key - The top-level key to load state_dict from the file.

pub fn with_debug_print(self) -> Self

Sets printing debug information on.

Trait Implementations

impl Clone for LoadArgs

fn clone(&self) -> LoadArgs

Returns a copy of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for LoadArgs

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

Formats the value using the given formatter.

impl From<&str> for LoadArgs

fn from(val: &str) -> Self

Converts to this type from the input type.

impl From<PathBuf> for LoadArgs

fn from(val: PathBuf) -> Self

Converts to this type from the input type.

impl From<String> for LoadArgs

fn from(val: String) -> Self

Converts to this type from the input type.