[ANN] PyTorch bindings

We are very pleased to announce the first release of our PyTorch bindings for ocaml. These bindings provide a NumPy like tensor library with GPU acceleration and support for automatic differentiation.

The torch package can be installed via opam. More details can be found in the github repo, including a tutorial training some neural networks on the MNIST dataset, various deep-learning examples: Generative Adverserial Networks, Neural Style Transfer, state of the art computer vision models on CIFAR-10, etc. Pre-trained weights for some recent computer vision models are also provided.

Any feedback is very welcome. Hopefully we will write a couple new tutorials to describe the current examples and add more examples in the next few weeks.

(note that the tensorflow package providing TensorFlow bindings also has been updated this week and is now up to speed with the github repo)


Super useful!
Thanks a lot.
I will very probably play with it in the beginning of next year.

related: Is it possible to use Machine Learning and Deep Learning frameworks through OCaml?


How does one load a pytorch pickle checkpoint in your binding?

Thanks for making them!

There is no mechanism to process pickle checkpoints on the OCaml side at the moment (this would essentially require a Python runtime). Instead there are two ways to use Python defined weights/models in OCaml.

  • It is possible to export the weight from Python to a .npz file and then convert them via the tensor_tools executable. The weights can then be imported by an OCaml defined model. This is how all the pre-trained model weights were created, see this resnet expample.
  • Another possibility is to export a JIT model from Python and import it on the OCaml side. There are some details in this tutorial.

The main difference is that in the first case, the model is defined via some OCaml code, whereas in the second case the OCaml side imports some “abstract” model.

1 Like