[ONNXRuntimeError] : 10 : INVALID_GRAPH 加载模型报错

项目场景：

Pytorch模型转ONNX模型，可以成功导出，但使用onnxruntime加载模型时出现如下错误

InvalidGraph: [ONNXRuntimeError] : 10 : INVALID_GRAPH : Load model from T.onnx failed:This is an invalid model. Type Error: Type ‘tensor(bool)’ of input parameter (8) of operator (ScatterND) in node (ScatterND_15) is invalid.

---

问题描述：

import torch
import torch.nn as nn
import onnxruntime
from torch.onnx import export

class Preprocess(nn.Module):
    def __init__(self):
        super().__init__()
        self.max = 1000
        self.min = -44

    def forward(self, inputs):
        inputs[inputs>self.max] = self.max
        inputs[inputs<self.min] = self.min
        return inputs
        
x = torch.randint(-1024,3071,(1,1,28,28))
model = Preprocess()
model.eval()

export(
    model,
    x,
    "test.onnx",
    input_names=["input"],
    output_names=["output"],
    opset_version=11,
)

session = onnxruntime.InferenceSession("test.onnx")

---

原因分析：

在pytorch的github中可以查到同样的问题#34054

---

解决方案：

ONNX的@shubhambhokare1修复了这个Issue，详情可见#45584

具体的操作为：先生成mask，再使用torch.masked_fill()操作。而不是利用索引直接对输入的tensor进行赋值

class MaskHigh(nn.Module):
    def __init__(self, val):
        super().__init__()
        self.val = val

    def forward(self, inputs):
        x = inputs.clone()
        mask = x > self.val
        output = torch.masked_fill(inputs, mask, self.val)
        return output


class MaskLow(nn.Module):
    def __init__(self, val):
        super().__init__()
        self.val = val

    def forward(self, inputs):
        x = inputs.clone()
        mask = x < self.val
        output = torch.masked_fill(inputs, mask, self.val)
        return output


class Clip(nn.Module):
    def __init__(self):
        super().__init__()
        self.high = MaskHigh(1300)
        self.low = MaskLow(-44)

    def forward(self, inputs):
        output = self.high(inputs)
        output = self.low(output)
        return output

可以利用Netron来可视化前后两种方式生成的计算图

索引赋值

Mask填充