ocaml-torch

ocaml-torch provides OCaml bindings for the PyTorch tensor library. This brings to OCaml NumPy-like tensor computations with GPU acceleration and tape-based automatic differentiation.

These bindings use the PyTorch C++ API and are mostly automatically generated. The current GitHub tip corresponds to PyTorch v2.3.

Installation with Libtorch

Torch depends on libtorch, so when you install this package, it will try linking to libtorch depending on your environment variables. The code for discovering libtorch is in src/config/discover.ml. In order to change how torch binds to libtorch, you must uninstall and reinstall torch.

To install with any of these methods, after configuring your environment, you may either

• opam install torch, or
• build from source:

  git clone https://github.com/janestreet/torch.git
  cd torch
  make all

Option 1: OPAM switch (CPU only)

The opam libtorch package (an optional dependency) can be installed, and torch will automatically detect it and build with it. However, it might not suit your needs if you use any of these:

• Windows operating system,
• ARM processors, or
• GPUs.

Option 2: Conda

If you've installed libtorch via Conda, ensure that you are in the Conda environment with the CONDA_PREFIX set before installing.

Option 3: System Libraries

If you have libtorch installed as a system library (e.g. RPM), set LIBTORCH_USE_SYSTEM=1 before installing.

Option 4: Custom Libtorch Location

If you have downloaded libtorch somewhere, set LIBTORCH=/path/to/libtorch/ before installing.

Examples

Utop

ocaml-torch can be used in interactive mode via utop or ocaml-jupyter.

Here is a sample utop session:

Simple Script

To build a simple torch program, create a file example.ml:

open Torch

let () =
  let tensor = Tensor.randn [ 4; 2 ] in
  Tensor.print tensor

Then create a dune file with the following content:

(executables
  (names example)
  (libraries torch))

Run dune exec example.exe to compile the program and run it!

Alternatively you can first compile the code via dune build example.exe then run the executable _build/default/example.exe (note that building the bytecode target example.bc may not work on macos).

Demos

• MNIST tutorial.
• Finetuning a ResNet-18 model.
• Generative Adversarial Networks.
• Running some Python model.
• ResNet examples on CIFAR-10.
• Character-level RNN
• Neural Style Transfer
• Reinforcement Learning

Some more advanced applications from external repos:

• An OCaml port of mini-dalle by Arulselvan Madhavan.
• Natural Language Processing models based on BERT can be found in the ocaml-bert repo.

Models and Weights

Various pre-trained computer vision models are implemented in the vision library. The weight files can be downloaded at the following links:

• ResNet-18 weights.
• ResNet-34 weights.
• ResNet-50 weights.
• ResNet-101 weights.
• ResNet-152 weights.
• DenseNet-121 weights.
• DenseNet-161 weights.
• DenseNet-169 weights.
• SqueezeNet 1.0 weights.
• SqueezeNet 1.1 weights.
• VGG-13 weights.
• VGG-16 weights.
• AlexNet weights.
• Inception-v3 weights.
• MobileNet-v2 weights.
• EfficientNet b0 weights, b1 weights, b2 weights, b3 weights, b4 weights.

Running the pre-trained models on sample images can the easily be done via:

dune exec examples/pretrained/predict.exe path/to/resnet18.ot images/tiger.jpg

Internals

ocaml-torch uses extensive code generation to produce bindings to thousands of torch C++ functions. Read internals.md for details.

Acknowledgements

Many thanks to @LaurentMazare for the original work of ocaml-torch.