Question

我正在尝试将pytorch中的简单模型（涉及conv和gru层）转换为onnx模型，并将其加载到Caffe。如果我使用训练有素的模型，则转换和Caffe加载效果很好。但是，我想删除模型的最后一层。尝试将其加载到Caffe时，出现错误“ RuntimeError：推断的形状和现有形状的等级不同：（0）与（2）”。

我描述的过程是：

训练完整模型
将其保存到.pth文件
创建一个没有最后一层的新模型对象
将state_dict加载到.pth中（使用strict = False，如果尝试加载时不这样做，则pytorch抱怨缺少最后一层）。
将其导出到onnx
将onnx模型加载到Caffe

我使用onnx.checker.check_model（model）函数检查是否一切正常，但到目前为止，我还没有得到任何输出。然后，执行

时程序崩溃

prepared_backend = onnx_caffe2_backend.prepare(model)

出现上述错误。

让我为这个错误困扰的是我不确定哪个层在抱怨。

编辑：

试图创建没有最后一层的模型，并立即将其导出，而无需进行培训。同样的问题。认为状态dict的加载方式会是一个问题，但看起来却不是这样。

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, :, :]

将pytorch模型转换为onnx

0 个答案: