PyTorch2ONNX2TensorRT 踩坑日誌

PyTorch2ONNX2TensorRT 踩坑日誌

從“用PyTorch寫的網絡，通過ONNX，使用TensorRT序列化，最終完成模型加速”的全流程踩坑日誌。

---

2019/12/07 初版
2019/12/17 更新AdaptivePooling, 找BUG思路
2019/12/27 添加AdaptivePooling示例
2020/01/01 添加VGG16示例鏈接

---

實驗環境

ONNX可以不用安裝，對ONNX2TRT沒有影響，推薦使用anaconda管理包。

• Ubuntu 16.04
• RTX2080TI, Driver Version: 410.79
• CUDA 10.0
• cudnn 7.6.3 (經測低版本如7.5.0無影響)
• pycuda 2019.1.2
• pytorch 1.3.1
• torchvision 0.4.2
• tensorrt 6.0.1.5
• python 3.6.9
  • 經測ONNX無法使用，建議使用python 3.7.x
  • onnx 1.6.0
  • protobuf 3.9.2 (需要降級到3.9.x，不然onnx會報libprotobuf.so.20的錯)

---

1. RuntimeError: ONNX export failed: Couldn’t export operator aten::upsample_bilinear2d

無法解決，ONNX2TensorRT報錯，待TensorRT後續版本支持，見後文替代方法#4

近似地，應該與警告信息 UserWarning: ONNX export failed on upsample_bilinear2d because align_corners == True not supported 相關聯。

原因

轉換ONNX使用的版本較低，PyTorch.ONNX不支持。另外，參考源碼， torch.onnx.export 默認使用 opset_version=9。

解決辦法

警告信息已經完整說明，ONNX's Upsample/Resize operator did not match Pytorch's Interpolation until opset 11.，因此將ONNX的導出代碼中規定其版本，具體如下：

import torch
torch.onnx.export(model, ..., opset_version=11)

較完整報錯信息

輸出的個人信息就被我隱去了，也爲了報錯、警告的簡潔，所以這裏叫做“較完整”，此說明後續不再贅述。

UserWarning: You are trying to export the model with onnx:Upsample for ONNX opset version 9. This operator might cause results to not match the expected results by PyTorch.
ONNX's Upsample/Resize operator did not match Pytorch's Interpolation until opset 11. Attributes to determine how to transform the input were added in onnx:Resize in opset 11 to support Pytorch's behavior (like coordinate_transformation_mode and nearest_mode).
We recommend using opset 11 and above for models using this operator. 

UserWarning: ONNX export failed on upsample_bilinear2d because align_corners == True not supported

RuntimeError: ONNX export failed: Couldn't export operator aten::upsample_bilinear2d

---

2. RuntimeError: ONNX export failed: Couldn’t export operator aten::adaptive_avg_pool2d

無法解決，ONNX2TensorRT報錯，待TensorRT後續版本支持，見後文替代方法#5

類似錯誤 aten::adaptive_avg_pool*d：onnx#63, pytorch#14395, discuss.pytorch#30204

原因

因爲PyTorch的網絡中用了 torch.nn.AdaptiveAvgPool2d ，個人感覺，ONNX沒有 avg_pool2d 操作，見ONNX Operator，所以PyTorch.ONNX就會報錯 aten::adaptive_avg_pool2d 無法轉換。

解決辦法

參考pytorch#14395添加額外Option，如下：

import torch
torch.onnx.export(model, ..., operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK)

該方法只是阻止ONNX替換PyTorch的OP、而是使用ATen的OP替換，PyTorch2ONNX能通，但ONNX2TRT卻不能通，原因是ONNX phaser識別不到非ONNX的OP。

較完整報錯信息

UserWarning: ONNX export failed on adaptive_avg_pool2d because output size that are not factor of input size not supported

RuntimeError: ONNX export failed: Couldn't export operator aten::adaptive_avg_pool2d

---

3. Error: In node 2 (importGather): UNSUPPORTED_NODE: Assertion failed: !(data->getType() == nvinfer1::DataType::kINT32 && nbDims == 1) && “Cannot perform gather on a shape tensor!”

原因

"Cannot perform gather on a shape tensor!"，網絡內部使用x_size = x.size()[1:]等類似操作，TensorRT在trace的時候，會被解析成一個shape layer的輸出，獲得一個shape tensor，用Netron工具可視化就可以發現，對應的node 2實際上是個Constant node，與預期不符。

解決辦法

不使用該操作，另一種解法來自onnx-tensorrt#192

x_size = torch.tensor(x.shape)[1:]

---

4. Error: In node 1 (importUpsample): UNSUPPORTED_NODE: Assertion failed: (nbDims >= 1) && (nbDims <= 3)

使用Netron工具可視化模型，找到對應的node 1，就可以發現對應的是F.interpolate(x, size=(128, 128), mode='bilinear', align_corners=False)操作。

原因

目前ONNX2TRT的轉換過程中，貌似不支持F.interpolate的bilinear模式，只支持linear和nearest。

解決辦法

將所有的bilinear模式替換爲nearest模式。

---

5. 使用AvgPooling替換AdaptivePooling

針對2. RuntimeError: ONNX export failed: Couldn't export operator aten::adaptive_avg_pool2d問題，使用AvgPooling替換AdaptivePooling。因爲ONNX支持AvgPooling，PyTorch2ONNX、ONNX2TRT流程能夠跑通。

原因

目前PyTorch2ONNX流程中，ONNX並不支持AdaptivePooling操作，該操作僅存於PyTorch中。

解決方法

參考[開發技巧]·AdaptivePooling與Max/AvgPooling相互轉換一文、PyTorch官方文檔可知，AdaptivePooling可通過輸入大小input_size自適應控制輸出大小output_size，而一般的AvgPooling/MaxPooling則是通過kernel_size、stride、padding來計算output_size，公式如下：

$\mathbf{output\_size} = ceil(( \mathbf{input\_size} + 2 * \mathbf{padding} - \mathbf{kernel\_size}) / \mathbf{stride})+1$

因此通過input_size、output_size反推kernel_size、stride、padding，參考官方源碼將padding設爲0，那麼可推出去kernel_size、stride：

$\mathbf{stride} = floor(\mathbf{input\_size} / \mathbf{output\_size})$
$\mathbf{kernel\_size} = \mathbf{input\_size}- (\mathbf{output\_size}-1) * \mathbf{stride}$

示例

例如，PyTorch網絡的某一層含有nn.AdaptiveAvgPool2d(output_size=(14,14))，它的output_size爲(14, 14)，該層的輸入特徵圖大小爲10*128*128，那麼輸出的特徵圖大小爲10*14*14，那麼帶入公式，可計算出nn.AvgPool2d(kernel_size, stride)的stride=(int(128/14), int(128/14)), kernel_size=((128-(14-1)*stride, (128-(14-1)*stride)，驗證如下：

import torch
from   torch import nn

input = torch.randn(10, 36, 36)
AAVP = nn.AdaptiveAvgPool2d(output_size=(12,12))
AVP  = nn.AvgPool2d(kernel_size=(3,3), stride=(3,3))

output_AAVP = AAVP(input)
output_AVP  = AVP(input)

---

6. PyTorch2ONNX、ONNX2TRT到底哪裏出了問題？

下面是遇到無法解決的問題後該找誰問的一個思路：

PyTorch2ONNX是調用的PyTorch內部的轉換腳本，所以PyTorch2ONNX出了問題，需要去PyTorch那裏的issue尋找解決辦法；ONNX2TRT是使用ONNX自己寫的轉換腳本onnx-tensorrt，同理如果ONNX2TRT出了問題，就需要到它的那裏找解決辦法；在產生好TRT模型後，使用TRT模型進行推理出問題了，那就要去TRT那裏問了，有GitHub和官方論壇可以使用。

那怎麼讓報錯暴露出來呢，下面是一些辦法。

解決方法

按下列方法多半能找到問題所在。

1. PyTorch2ONNX

1. 更新PyTorch到最新版，一般最新版中ONNX的OP支持應該會更多；
2. 按下列代碼將日誌等級調到最高，逐一分析。

import torch
torch.onnx.export(..., verbose=True, ...)

2. 檢測ONNX模型

下載Netron可視化自己的ONNX模型，分析是否與PyTorch模型一致，或者與自己想造的模型一致。多看看resize、shape、permute操作，ONNX對這些需要對tensor切片的操作不是很支持。

3. ONNX2TRT

1. 更新onnx-tensorrt庫，也就是libnvonnxparser.so。下面貼一段TRT的安裝步驟：
   1. 安裝TRT.
   2. 編譯onnx-tensorrt.
   3. 將libnvonnxparser.so移到TRT的lib文件夾中.
2. 按下列代碼將日誌等級調到最高，逐一分析。

import tensorrt as trt
TRT_LOGGER = trt.Logger(trt.Logger.VERBOSE)

---

最終解決辦法

放棄ONNX2TRT吧，PyTorch與ONNX與TRT的版本難以互相支持，在版本的迭代中任意節點不支持了，整個鏈路就會斷掉，另外TRT是閉源的項目，你完全不知道ONNX2TRT的過程中出了哪些問題，就算有堆棧信息，也不可能根據信息去trace它的錯誤。所以，直接使用TRT提供的api直接構建網絡，是最複雜、也是最簡單直接的方法。

Pytorch 2 TRT python API

使用TRT提供的python接口，構建網絡，整個流程十分簡單，大家可以看看TRT提供的示例<TRT_root>/samples/python/network_api_pytorch_mnist/sample.py，與之對照的是<TRT_root>/samples/python/network_api_pytorch_mnist/model.py：

def populate_network(network, weights):
    # Configure the network layers based on the weights provided.
    input_tensor = network.add_input(name=ModelData.INPUT_NAME, dtype=ModelData.DTYPE, shape=ModelData.INPUT_SHAPE)
    """
    TRT python API
    """
    network.mark_output(tensor=fc2.get_output(0))

你只需要把這個populate_network寫出來就好了，weights就是網絡的權重了，由torch.load()得到，是不是超級簡單啊。想使用PyTorch的F.interpolate的bilinear模式？TRT提供！下篇日誌將會記錄“如何使用TRT python API搭建簡單的VGG16網絡”，我再也不想用ONNX2TRT了。