pytorch模型转onnx,报错kind_.is_prim() INTERNAL ASSERT FAILED at /pytorch/torch/csrc/jit/ir.cpp:922
在将pytorch模型转为onnx模型时遇到一个报错如下:
这个报错指向范围很宽泛,代码有几种不同的错误都会报这个错误,我在本文里面列出我遇到的这个报错但是实际上是不同类型的代码问题以及解决方法,目前我只遇到两个报错相同但实际上代码错误不同的问题,未来遇到还会补充。
1.创建tensor的问题
我遇到的第一个这个报错是在使用torch.Tensor()方法创建新的tensor,最终的解决方法是用torch.tensor()代替torch.Tensor()来创建新的tensor,所以推测在转换onnx模型的时候不支持使用torch.Tensor()来创建tensor。
2.计算过程中的赋值操作没有带‘=’
我在使用torch.topk()来进行操作的时候出现了这个报错,刚开始我使用torch.topk的源代码是torch.topk(input, k=5, dim=-1, out=(outpu1,outpu2)),这样使用torch.topk()会直接把结果赋值给output1和output2,但是在转换onnx模型时就出现了这个报错,最终的解决方法是把源代码改为output1,output2=torch.topk(input,k=5,dim=-1),因此可以推测转换onnx模型的时候必须使用带等号这种形式的赋值操作。