«RuntimeError: Предполагаемая форма и существующая форма различаются по рангу: (0) против (2)» при преобразовании модели pytorch в onnx - PullRequest
0 голосов
/ 04 апреля 2019

Я пытаюсь преобразовать простую модель (включающую слои conv и gru) в pytorch в модель onnx и загрузить ее в Caffe. Если я использую полностью обученную модель, преобразование и загрузка Caffe будут работать нормально. Тем не менее, я хочу отбросить последний слой модели. Когда я делаю это и пытаюсь загрузить его в Caffe, я получаю сообщение об ошибке «RuntimeError: Предполагаемая форма и существующая форма различаются по рангу: (0) против (2)».

Процедура, которую я описываю:

  • Поезд полной модели
  • Сохраните его в файл .pth
  • Создание нового объекта модели без конечного слоя
  • Загрузите state_dict в .pth (со строгим = False, если при попытке загрузить его pytorch жалуется на отсутствие последнего слоя).
  • Экспорт в onnx
  • Загрузить модель onnx в Caffe

Я использую функцию onnx.checker.check_model (model), чтобы проверить, все ли в порядке, но пока я не получаю никаких выходных данных. Затем программа аварийно завершает работу при выполнении

prepared_backend = onnx_caffe2_backend.prepare(model)

где я получаю ошибку, указанную выше.

Что меня беспокоит в этой ошибке, так это то, что я не уверен, какой слой жалуется.

EDIT:

Попытался создать модель без конечного слоя и экспортировать ее сразу, без обучения. Та же проблема. Подумал, что проблема в том, как загружается государственный диктат, но не выглядит так.

EDIT2:

Вот код пересылки:

    x = self.embeddings(x)
    x = x.permute(0, 3, 1, 2)

    x = self.conv1(x).squeeze(3).permute(2, 0, 1)

    outputs, hidden = self.gru1(x)

    if self.with_output_layer:
        output = outputs.permute(1, 0, 2)[:, 25, :]
        return self.out_layer(output)
    else:
        # sum last state from forward and backward  direction
        return hidden[3, :, :] + hidden[2, :, :]
...