Passing Python torch.nn.Module to c++ torch::nn::Module

Hi guys, I want to do forward pass of my pytorch models that is defined in python with C++ and pybind. I followed the tutorial on cuda/C++ extension here. With that I was able to:

  1. Wrote ext.cpp where the C++ function for forward pass resides


#include <torch/extension.h> // One-stop header.

at::Tensor extforward(torch::nn::Module *model, torch::Tensor x)
for(const auto submod : model->modules())
//Looping through the network’s every submodule and applying forward to input
x = submod->astorch::nn::AnyModule()->forward(x);
return x;

m.def(“forw”, &extforward, “CPP forward”);

  1. Build a that compiles the C++ function and make an API sucessfully

from setuptools import setup, Extension
from torch.utils import cpp_extension

libname = ‘ext’
ext_modules=[cpp_extension.CppExtension(libname, [‘ext.cpp’])],
cmdclass={‘build_ext’: cpp_extension.BuildExtension})


  1. Wrote that builds a simple network and try to do forward pass with c++ function

import torch
import torch.nn as nn
import torch.nn.functional as F
import ext as ext

class Net(nn.Module):

def __init__(self):
    super(Net, self).__init__()
    # 1 input image channel, 6 output channels, 3x3 square convolution
    # kernel
    self.conv1 = nn.Conv2d(1, 6, 3)
    self.conv2 = nn.Conv2d(6, 16, 3)
    # an affine operation: y = Wx + b
    self.fc1 = nn.Linear(16 * 6 * 6, 120)  # 6*6 from image dimension
    self.fc2 = nn.Linear(120, 84)
    self.fc3 = nn.Linear(84, 10)

def forward(self, x):
    # Max pooling over a (2, 2) window
    x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
    # If the size is a square you can only specify a single number
    x = F.max_pool2d(F.relu(self.conv2(x)), 2)
    x = x.view(-1, self.num_flat_features(x))
    x = F.relu(self.fc1(x))
    x = F.relu(self.fc2(x))
    x = self.fc3(x)
    return x

def num_flat_features(self, x):
    size = x.size()[1:]  # all dimensions except the batch dimension
    num_features = 1
    for s in size:
        num_features *= s
    return num_features

if name == “main”:
net = Net()
x = torch.randn(1, 1, 32, 32)

Running Instruction:

  1. python --install
  2. python


  1. The C++ API does not recognize python’s custom network’s type as seen below:

Other Details:

  1. System: Ubuntu 18.04
  2. Environment: virtual env with python3.6
  3. Pytorch version: 1.5.1 stable

I know we can use torch jit script, but I feel like torch::nn::Module has much capability that torch::script::module do not have. Would there be anyways to sort of extract all the information from torch.nn.Module into torch::nn::Module? Thanks in advance for all your feedback guys! :slight_smile:


no. torch.nn.Module and torch::nn::Module are completely independent. While we aim for parity between the APIs, you cannot convert from one to another.
The deeper reason is that it would give up quite bit of “Pythonic” to move the modular interface into C++.
So indeed, I think JIT might be the only way to use the same module on both sides of C++ / Python (unless you want to call back into Python from C++).

Best regards



@tom Could you please clarify, is it still true that we cannot convert from one API to another as of now (a couple years later)?

Best regards

“cannot convert” is decidedly too strong.
A lot of progress has been made.
You still need to do a bit of a dance because the Python API for nn.Module itself does not translate to C++ in a completely “just add semicolons and a few auto typing declarations”-way (e.g. forward is actually tricky because of the signatures).

Best regards


1 Like

Some details on how to achieve this would be helpful.