Integrating Pre-Trained PyTorch Models into Your R Package

This note uses EvoNN to illustrate how an R package can call a trained PyTorch model through reticulate. The examples assume familiarity with R packages and PyTorch; model architecture and tree preprocessing must match the training code.

Manage Python dependencies

The EvoNN package layout below includes a version list, inst/pkglist.csv. Pinning tested versions helps reproduce model predictions.

The table lists dependencies for PyTorch and PyTorch Geometric:

Organize models and Python scripts

The /inst directory also stores pre-trained neural network models (weights.pt) and Python scripts (import.py and function.py) that contain the necessary libraries, functions defining the neural network architecture, and data loading mechanisms.

Expand the folders below to inspect EvoNN’s package structure.

Prepare Python source files

We can divide our Python scripts into two parts. The first part, import.py, contains the necessary Python libraries to load the pre-trained model and perform the neural network estimation. The second part, function.py, contains the Python function that performs the neural network estimation.

For example, in the import.py file you can write this:

# import.py
import torch
import torch_geometric
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

And in the function.py you can write a model-specific inference function:

# function.py
# Model-specific preprocessing and output conversion are placeholders.
def py_function(py_tree, weights_path):
    def create_dataset(tree):
        # Define the dataset creation process
        return dataset

    py_dataset = create_dataset(py_tree)

    # Define the neural network architecture
    class Net(nn.Module):
        def __init__(self):
            super(Net, self).__init__()
            self.conv1 = GCNConv(16, 16)
            self.conv2 = GCNConv(16, 16)
            self.fc = nn.Linear(16, 1)

        def forward(self, data):
            x, edge_index = data.x, data.edge_index
            x = F.relu(self.conv1(x, edge_index))
            x = F.relu(self.conv2(x, edge_index))
            x = F.dropout(x, training=self.training)
            x = self.fc(x)
            return x

    # Load the pre-trained model
    model = Net()
    model.load_state_dict(torch.load(weights_path, map_location="cpu", weights_only=True))
    model.eval()

    # Perform neural network estimation
    with torch.inference_mode():
        out = model(py_dataset)
    # Convert to a data dict of numpy arrays
    out = convert_to_numpy(out)
    return out

For a self-contained Python module, put the imports at the top of function.py. The example architecture, preprocessing and output conversion need to be replaced with those used to train your model.

Declare Python dependencies

With reticulate 1.41 or later, declare dependencies in .onLoad() using py_require(). Reticulate resolves them when Python is first used. Loading the R package should not install packages into, or replace, a user’s chosen environment. See the reticulate package guide.

# R/zzz.R
.onLoad <- function(libname, pkgname) {
  reticulate::py_require(c("torch", "torch-geometric"))
}

Declare reticulate (>= 1.41) in Imports. Add tested version constraints to the Python requirements for a released model. Users who select their own Python environment must install compatible dependencies there.

Call the model from R

Import the module when the estimation function is called. Resolve the weights through system.file() so prediction does not depend on the working directory. In this sketch, prepare_tree() is the package’s model-specific preprocessing function.

nn_estimate <- function(phylo_tree) {
  model_dir <- system.file("model", package = "EvoNN", mustWork = TRUE)
  weights <- system.file("model", "weights.pt", package = "EvoNN", mustWork = TRUE)
  model <- reticulate::import_from_path("function", path = model_dir)
  model$py_function(prepare_tree(phylo_tree), weights)
}

Reticulate converts supported R objects to Python and converts the returned result back to R. The input shapes, feature ordering and tensor types must agree with the trained model.

Check the package

Run R CMD check, test prediction against known outputs, and check that loading the package works without initializing Python. Keep examples that require model weights or Python dependencies separate from examples that only exercise R code.

For the complete EvoNN implementation, see the repository:

EvoNNView source on GitHub