This Rust crate provides low-level bindings to the C++ PyTorch (libtorch) API, enabling Rust developers to leverage PyTorch's deep learning capabilities. It aims to mirror the C++ API closely, serving as a foundation for higher-level Rust abstractions.

How It Works

The tch-rs crate acts as a thin wrapper around libtorch, the C++ backend of PyTorch. It utilizes code generation derived from ocaml-torch to create bindings for the C API of libtorch. This approach ensures a direct mapping to PyTorch's core functionalities, allowing for efficient tensor operations, automatic differentiation, and neural network construction within a Rust environment.

Quick Start & Requirements

• Installation: Requires libtorch version v2.7.0. Can be linked to a system-wide installation, a Python PyTorch installation (LIBTORCH_USE_PYTORCH=1), or downloaded via the download-libtorch feature. CUDA support requires TORCH_CUDA_VERSION to be set (e.g., cu117).
• Manual Libtorch: Set LIBTORCH environment variable to the libtorch directory.
• Static Linking: Use LIBTORCH_STATIC=1 and compile libtorch manually.
• Windows: Recommended to use MSVC Rust toolchain due to PyTorch compatibility.
• Docs: docs.rs

Highlighted Details

• Provides tensor operations, neural network modules (nn), optimizers, and utilities for vision tasks (e.g., MNIST, ResNet).
• Supports loading pre-trained models, including those exported via safetensors.
• Includes examples for gradient descent, neural network training, and inference with pre-trained models.

Maintenance & Community

• Active development with examples covering various deep learning tasks.
• Community discussions may occur on GitHub issues.

Licensing & Compatibility

• Dual-licensed under MIT and Apache 2.0. Permissive for commercial use and integration with closed-source projects.

Limitations & Caveats

• Compilation can be slow due to potential torch-sys rebuilds; environment variables like LIBTORCH need to be correctly configured for IDEs like rust-analyzer.
• Windows users should be mindful of MSVC vs. MinGW compatibility issues with PyTorch.
• Error handling for shared library loading requires careful configuration of LD_LIBRARY_PATH or DYLD_LIBRARY_PATH.