tch-rs

To avoid Python pickle dependencies and enable zero-copy loading, use the safetensors format.

1. Export in PyTorch: Use the safetensors Python library to save the state_dict() with a .safetensors suffix.
2. Import in tch: Use vs.load("filename.safetensors")? to load the weights into your VarStore.

Note: The filename must have the .safetensors suffix for tch to decode it correctly.

use anyhow::Result;
use tch::{
	Device,
	Kind,
	nn::VarStore,
	vision::{
		imagenet,
		resnet::resnet18,
	}
};

fn main() -> Result<()> {
	// Create the model and load the pre-trained weights
	let mut vs = VarStore::new(Device::cuda_if_available());
	let model = resnet18(&vs.root(), 1000);
	vs.load("resnet18.safetensors")?;
	
	// Load the image file and resize it to the usual imagenet dimension of 224x224.
	let image = imagenet::load_image_and_resize224("dog.jpg")?
		.to_device(vs.device());

	// Apply the forward pass of the model to get the logits
	let output = image
		.unsqueeze(0)
		.apply_t(&model, false)
		.softmax(-1, Kind::Float);
	
	// Print the top 5 categories for this image.
    for (probability, class) in imagenet::top(&output, 5).iter() {
        println!("{:50} {:5.2}%", class, 100.0 * probability)
    }
    
    Ok()
}

To train models, use nn::VarStore to manage variables and their initializations. Optimization is performed using optimizers like nn::Sgd. The typical workflow involves:

1. Defining a module using nn::Path to create variables.
2. Initializing a VarStore on a specific Device.
3. Building an optimizer (e.g., nn::Sgd) from the VarStore.
4. In a training loop: performing a forward pass, computing loss, and calling opt.backward_step(&loss) to compute gradients and update variables.

use tch::nn::{Module, OptimizerConfig};
use tch::{kind, nn, Device, Tensor};

fn my_module(p: nn::Path, dim: i64) -> impl nn::Module {
    let x1 = p.zeros("x1", &[dim]);
    let x2 = p.zeros("x2", &[dim]);
    nn::func(move |xs| xs * &x1 + xs.exp() * &x2)
}

fn gradient_descent() {
    let vs = nn::VarStore::new(Device::Cpu);
    let my_module = my_module(vs.root(), 7);
    let mut opt = nn::Sgd::default().build(&vs, 1e-2).unwrap();
    for _idx in 1..50 {
        // Dummy mini-batches made of zeros.
        let xs = Tensor::zeros(&[7], kind::FLOAT_CPU);
        let ys = Tensor::zeros(&[7], kind::FLOAT_CPU);
        let loss = (my_module.forward(&xs) - ys).pow_tensor_scalar(2).sum(kind::Kind::Float);
        opt.backward_step(&loss);
    }
}

You can download libtorch from the PyTorch website and point the build script to it using the LIBTORCH environment variable.

Linux and macOS

Set the LIBTORCH variable in your .bashrc or equivalent to the path where you extracted libtorch.

If you need to specify header and library paths separately:

• LIBTORCH_INCLUDE must contain the include directory.
• LIBTORCH_LIB must contain the lib directory.

Windows

1. Create a LIBTORCH system environment variable pointing to the unzipped directory (e.g., X:\path\to\libtorch).
2. Append the lib subdirectory to your Path variable (e.g., X:\path\to\libtorch\lib).

Alternatively, use PowerShell for a temporary session:

$Env:LIBTORCH = "X:\path\to\libtorch"
$Env:Path += ";X:\path\to\libtorch\lib"

export LIBTORCH=/path/to/libtorch

# Optional: separate include/lib paths
export LIBTORCH_INCLUDE=/path/to/libtorch/
export LIBTORCH_LIB=/path/to/libtorch/

You can run pre-trained computer vision models (like ResNet) by loading weights into a VarStore. For example, using the pretrained-models example:

cargo run --example pretrained-models -- resnet18.ot tiger.jpg

In code, load the model architecture and then use vs.load(weight_file)? to populate the parameters.

    // First the image is loaded and resized to 224x224.
    let image = imagenet::load_image_and_resize(image_file)?;

    // A variable store is created to hold the model parameters.
    let vs = tch::nn::VarStore::new(tch::Device::Cpu);

    // Then the model is built on this variable store, and the weights are loaded.
    let resnet18 = tch::vision::resnet::resnet18(vs.root(), imagenet::CLASS_COUNT);
    vs.load(weight_file)?;

    // Apply the forward pass of the model to get the logits and convert them
    // to probabilities via a softmax.
    let output = resnet18
        .forward_t(&image.unsqueeze(0), /*train=*/ false)
        .softmax(-1);

    // Finally print the top 5 categories and their associated probabilities.
    for (probability, class) in imagenet::top(&output, 5).iter() {
        println!("{:50} {:5.2}%", class, 100.0 * probability)
    }

Use nn::seq() to compose layers like nn::linear. You can add functional layers using .add_fn(). For training, use specialized loss functions like .cross_entropy_for_logits() and optimizers like nn::Adam.

use anyhow::Result;
use tch::{nn, nn::Module, nn::OptimizerConfig, Device};

const IMAGE_DIM: i64 = 784;
const HIDDEN_NODES: i64 = 128;
const LABELS: i64 = 10;

fn net(vs: &nn::Path) -> impl Module {
    nn::seq()
        .add(nn::linear(
            vs / "layer1",
            IMAGE_DIM,
            HIDDEN_NODES,
            Default::default(),
        ))
        .add_fn(|xs| xs.relu())
        .add(nn::linear(vs, HIDDEN_NODES, LABELS, Default::default()))
}

pub fn run() -> Result<()> {
    let m = tch::vision::mnist::load_dir("data")?;
    let vs = nn::VarStore::new(Device::Cpu);
    let net = net(&vs.root());
    let mut opt = nn::Adam::default().build(&vs, 1e-3)?;
    for epoch in 1..200 {
        let loss = net
            .forward(&m.train_images)
            .cross_entropy_for_logits(&m.train_labels);
        opt.backward_step(&loss);
        let test_accuracy = net
            .forward(&m.test_images)
            .accuracy_for_logits(&m.test_labels);
        println!(
            "epoch: {:4} train loss: {:8.5} test acc: {:5.2}%",
            epoch,
            f64::from(&loss),
            100. * f64::from(&test_accuracy),
        );
    }
    Ok(())
}