Using torch.load from Julia


Creation date: 2021-03-14

Author: Gregory Sech

Tags: pycall, pytorch

As part of an experiment we needed to solve a large quantity of linear inequalities described by the coefficient of a very sparse matrix.

These coefficients were calculated using PyTorch by a colleague and, after solving them within Python with an ad-hoc algorithm, I wanted to solve the system by using a LP solver via JuMP.

I'm a novice of Julia but recently, on the Beginners AMA featuring Dr. Katharine Hyatt and Dr. Rachel Kurchin, I've heard about an interesting package called PyCall.jl. This provides the ability to call Python modules directly from Julia programs.

The setup was quite easy, I just had to follow the instructions of PyCall.jl's README about using a different version of Python as I have installed PyTorch in a Anaconda's environment.

It was really easy, I just copied the path to the correct executable and before doing add PyCall I've set the PYTHON environment variable directly from the Julia REPL then entered Pkg REPL-mode by pressing ] and finally added PyCall. For a full description of the procedure everything is described in the README mentioned above.

After that loading data from PyTorch was pretty straightforward but let me walk you through it:

using SparseArrays, PyCall

@pyimport torch

data_dict = torch.load("<path to my data>", map_location=torch.device("cpu"))
C_rows = data_dict["ids"][0].numpy() .+ 1
C_cols = data_dict["ids"][1].numpy() .+ 1
C_values = data_dict["values"].numpy()
C = sparse(C_rows, C_cols, C_values)
  1. I've used PyCall to import the torch module from Python.

  2. Called torch.load as if I was in Python (with the only difference that Julia doesn't allow for single quotes as string delimiters) to load my data, which is a dictionary of tensors.

  3. By reading PyCall's README I've discovered that numpy's arrays are automatically converted into an interoperable array type in Julia. I just needed to add a +1 as Julia's indexing starts from 1 by default and I didn't really want to mess with custom indices.

  4. After I've got all my data loaded I just needed to create my sparse matrix C using sparse from the SparseArrays module of the Standard Library.

So today we learned that Julia's interoperability with Python is quite easy, even with non standard modules like PyTorch, and that using Numpy's arrays to move data between languages is really neat and painless!

Subscribe to RSS