экспорт модели pytorch через onnx в caffe2 - PullRequest
Новая
       51

экспорт модели pytorch через onnx в caffe2

0 голосов
/ 09 января 2019

Я пытаюсь обновить существующую модель pytorch 0.4 до версии 1.0 и пытаюсь использовать бэкэнд Caffe2 для запуска моделей в производстве на графическом процессоре.

Итак, я сделал следующее: 

# 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: Что если мои входные данные уже находятся на графическом процессоре? Как я могу передать эти данные в модель вместо того, чтобы переместить их в 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
  • cuda 9.0
...