Экспорт модели pytorch через onnx в caffe2
Я пытаюсь обновить мою существующую модель 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? Как я могу передать эти данные в модель вместо того, чтобы переместить их в CPU с numpy, а затем передать их исполнителю? Я попробовал следующее:
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? Итак, относительно этого момента. Если я позвоню prepare без опции GPU, вызов быстрый, но с указанием GPU onnx_caffe2_backend.prepare(model, device="CUDA:3") очень медленно
Моя система использует
- питон 3.6.8
- pytorch 1.0.0
- onnx 1.3.0
- Ubuntu 16.04
- куда 9.0