#ffi #deep-learning #machine-learning #pytorch

sys torch-sys

Low-level FFI bindings for the PyTorch C++ api (libtorch)

5 releases

✓ Uses Rust 2018 edition

new 0.0.5 Mar 19, 2019
0.0.4 Mar 7, 2019
0.0.3 Mar 3, 2019
0.0.2 Mar 2, 2019
0.0.1 Mar 2, 2019

#34 in Machine learning

Download history 19/week @ 2019-02-25 34/week @ 2019-03-04 17/week @ 2019-03-11

27 downloads per month
Used in 1 crate

Apache-2.0

141KB
1.5K SLoC

Rust 1K SLoC // 0.0% comments C++ 527 SLoC // 0.0% comments

tch-rs

Experimental rust bindings for PyTorch. The code generation part for the C api on top of libtorch comes from ocaml-torch.

Instructions

export LIBTORCH=/path/to/libtorch
export LD_LIBRARY_PATH=${LIBTORCH}/lib:$LD_LIBRARY_PATH
  • You should now be able to run some examples, e.g. cargo run --example basics.

Examples

The following code defines a simple model with one hidden layer.

struct Net {
    fc1: nn::Linear,
    fc2: nn::Linear,
}

impl Net {
    fn new(vs: &mut nn::VarStore) -> Net {
        let fc1 = nn::Linear::new(vs, IMAGE_DIM, HIDDEN_NODES);
        let fc2 = nn::Linear::new(vs, HIDDEN_NODES, LABELS);
        Net { fc1, fc2 }
    }
}

impl nn::Module for Net {
    fn forward(&self, xs: &Tensor) -> Tensor {
        xs.apply(&self.fc1).relu().apply(&self.fc2)
    }
}

This model can then be trained on the MNIST dataset with the following code.

fn main() {
    let m = tch::vision::mnist::load_dir(std::path::Path::new("data")).unwrap();
    let mut vs = nn::VarStore::new(Device::Cpu);
    let net = Net::new(&mut vs);
    let opt = nn::Optimizer::adam(&vs, 1e-3, Default::default());
    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),
        );
    }
}

More examples can be found in the examples directory. They can be run using the following command:

cargo run --example mnist_nn

Dependencies

~1.5MB
~28K SLoC

  • libc 0.2
  • links torch