Я пытаюсь преобразовать простую модель (включающую слои 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, :, :]