How to run a pytorch-onnx-caffe2 model on GPU?

[perrywu1989]

envs:
Pytorch of latest version(torch 1.0.0a0+058c12) from source.
|--- install build-in convert-onnx-to-caffe2 tool ---|
Onnx (onnx 1.3.0) from pip
conda python 3.5

First:
I use codes below to convert pytorch to onnx. Works!

model = Resnet18()
if torch.cuda.is_available():
    state_dict = torch.load(model_file)
    model.load_state_dict(state_dict)
    model.cuda()
    x = torch.randn(1, 3, 128, 128, requires_grad=True).cuda()
else:
    state_dict = torch.load(model_file, map_location='cpu')
    model.load_state_dict(state_dict)
    x = torch.randn(1, 3, 128, 128, requires_grad=True)
model.train(False)
torch_out = torch.onnx._export(model,  
                               x,  
                               "ocr_api.onnx",  
                               export_params=True)

After I got ONNX, I convert it to caffe2.Works!

Bash Command: onvert-onnx-to-caffe2 ocr_api.onnx --output predict_net.pb --init-net-output init_net.pb

After I got two .pb files, I run them on Caffe2. Works but in CPU model!

with open("init_net.pb",'rb') as f:
    init_net = f.read()
with open("predict_net.pb",'rb') as f:
    predict_net = f.read()
p = workspace.Predictor(init_net, predict_net)
img = np.random.randn(1, 3, 128, 128).astype(np.float32)
out = p.run([img])

Unfortunately, I try many ways to make them run on GPU, but not any works.
For examplse:

import numpy as np
import os, time
from caffe2.python import core, workspace, model_helper
from caffe2.proto import caffe2_pb2, caffe2_legacy_pb2

workspace.ResetWorkspace()
device_opts = core.DeviceOption(caffe2_pb2.CUDA, 0)

net_def = caffe2_pb2.NetDef()
with open('predict_net.pb', 'rb') as f:
    net_def.ParseFromString(f.read())
    net_def.device_option.CopyFrom(device_opts)
    workspace.CreateNet(net_def)
    # workspace.CreateNet(net_def.SerializeToString())

init_def = caffe2_pb2.NetDef()
with open('init_net.pb', 'rb') as f:
    init_def.ParseFromString(f.read())
    init_def.device_option.CopyFrom(device_opts)
    workspace.RunNetOnce(init_def)
    # workspace.RunNetOnce(init_def.SerializeToString())

name = net_def.name
output_name = net_def.external_output[-1]
input_name = net_def.external_input[0]


input_data = np.random.rand(1, 3, 32, 128).astype(np.float32) # NCHW

workspace.FeedBlob(input_name, input_data, device_opts) # workspace.FeedBlob('data', input_data, device_opts)

workspace.RunNet(name, 1)
results = workspace.FetchBlob(output_name)

it said:

Traceback (most recent call last):
  File "run_caffe2.py", line 78, in <module>
    workspace.CreateNet(net_def)
  File "/home/wfy/anaconda3/envs/caffe2/lib/python3.5/site-packages/caffe2/python/workspace.py", line 154, in CreateNet
    StringifyProto(net), overwrite,
  File "/home/wfy/anaconda3/envs/caffe2/lib/python3.5/site-packages/caffe2/python/workspace.py", line 180, in CallWithExceptionIntercept
    return func(*args, **kwargs)
RuntimeError: [enforce fail at operator.cc:46] blob != nullptr. op Conv: Encountered a non-existing input blob: 0

Even I use the codes in comment lines, the error messages are the same.

Is there anyone who knows how to solve this problem? Or how to make a pytorch-onnx-caffe2 model
run on GPU?
Thanks very much!

cc @houseroad @spandantiwari @lara-hdr @BowenBao @neginraoof

[mhubii]

Could you try this?

with open("init_net.pb",'rb') as f:
    init_net = f.read()
with open("predict_net.pb",'rb') as f:
    predict_net = f.read()

init_net.RunAllOnGPU() # Added this line.
predict_net.RunAllOnGPU() # Added this line.

p = workspace.Predictor(init_net, predict_net)
img = np.random.randn(1, 3, 128, 128).astype(np.float32)
out = p.run([img])

[perrywu1989]

@mhubii No, it still do not work.
the error is:

Traceback (most recent call last):
  File "run_caffe2.py", line 46, in <module>
    init_net.RunAllOnGPU()  # Added this line.
AttributeError: 'bytes' object has no attribute 'RunAllOnGPU'

[alexbuyval]

I have the same issue trying to run Caffe2 model on GPU:

    workspace.CreateNet(net_def)
  File "/home/alex/anaconda3/envs/pytorch_env/lib/python3.6/site-packages/caffe2/python/workspace.py", line 154, in CreateNet
    StringifyProto(net), overwrite,
  File "/home/alex/anaconda3/envs/pytorch_env/lib/python3.6/site-packages/caffe2/python/workspace.py", line 180, in CallWithExceptionIntercept
    return func(*args, **kwargs)
RuntimeError: [enforce fail at operator.cc:46] blob != nullptr. op Conv: Encountered a non-existing input blob: 0

@perrywu1989 Have you found a solution?

Thanks!

[satheeshkatipomu]

any found solution to this issue, Facing similar issue

[garymm]

The issue here is not with torch.onnx, so removing module: onnx label.