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 2.  With that I was able to:

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


ext.cpp
#include <torch/extension.h> // One-stop header.
#include 
#include 
#include 
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;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){
m.def(“forw”, &extforward, “CPP forward”);
}


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


setup.py
from setuptools import setup, Extension
from torch.utils import cpp_extension
libname = ‘ext’
setup(name=libname,
ext_modules=[cpp_extension.CppExtension(libname, [‘ext.cpp’])],
cmdclass={‘build_ext’: cpp_extension.BuildExtension})
print(cpp_extension.BuildExtension)


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


test.py
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)
ext.forw(net,x)

Running Instruction:

python setup.py --install
python test.py

Issue:

The C++ API does not recognize python’s custom network’s type as seen below:
Screenshot from 2020-07-18 01-38-131826×678 123 KB


Other Details:

System: Ubuntu 18.04
Environment: virtual env with python3.6
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!