我试图将现有的pytorch 0.4模型升级到1.0，并试图使用Caffe2后端在GPU上的生产环境中运行模型。
所以，我所做的如下：
# Export my model to ONNX
torch.onnx._export(model, args, "test.onnx", export_params=True)
import caffe2.python.onnx.backend as onnx_caffe2_backend
# Load the ONNX model from file.
model = onnx.load("test.onnx")
# We will run our model on the GPU with ID 3.
rep = onnx_caffe2_backend.prepare(model, device="CUDA:3")
outputs = rep.run(np.random.randn(1, 3, 128, 64).astype(np.float32))
现在，我对此有两个疑问：
1：如果我的输入数据已经位于GPU上怎么办？我如何将这些数据传递给模型，而不是使用numpy将其传递到CPU，然后传递给执行器？我尝试了以下方法：
args = torch.randn(1, 3, 128, 64, dtype=torch.float32).cuda(3)
print(args.dtype)
outputs = rep.run(args)
这将打印torch.float32。但是，我得到了错误：
if arr.dtype == np.dtype('float64'):
TypeError: data type not understood
我不确定为什么将数组解释为double数组。
2：我注意到对prepare的呼叫非常慢。因此，看来我的旧pytorch代码比在后端运行它要快。我将进行更详尽的时序比较，但这是导出模型并使其在带有pytorch / onnx / caffe的GPU上运行的正确方法吗？
因此，关于这一点。如果我在不使用GPU选项的情况下调用prepare，则调用速度很快，但是使用onnx_caffe2_backend.prepare(model, device="CUDA:3")指定GPU的速度非常慢。
我的系统正在使用