[BUG]Error Code 4: Internal Error (Internal error: plugin node ScatterND_12 requires 36 bytes of scratch space, but only 0 is available.

[MHGL]

bug描述
yolov5转换tensorrt过程中，如果网络中有存在以下现象，就会发生错误：

• 索引操作在等号左边
• 切片操作在等号左边

如何复现
复现步骤:

1. 最小代码

import torch

# init module
class MyModule(torch.nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        ...

    def forward(self, x):
        x[0] += 1
        return x

torch_model = MyModule()

# torch.onnx.export
torch.onnx.export(torch_model,
    torch.randn(1, 3, 256, 416),
    "./tmp.onnx",
    input_names=["inputs"],
    output_names=["outputs"],
    opset_version=11)

# onnx simplify
import os
import onnx
from onnxsim import simplify

onnx_file = os.path.join(os.getcwd(), "tmp.onnx")
model_op, check_ok = simplify(onnx_file, 
    check_n=3, 
    perform_optimization=True, 
    skip_fuse_bn=True,  
    skip_shape_inference=False, 
    input_shapes={"inputs": (1, 3, 256, 416)}, 
    skipped_optimizers=None, 
    )
onnx.save(model_op, "./tmp.onnx")

# onnx -> tensorrt
# !!!
# you should build tensorrt first
import tensorrt as trt

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
    with open(onnx_file, 'rb') as model:
        parser.parse(model.read())

    config = builder.create_builder_config()
    engine = builder.build_engine(network, config)

    with open("tmp.trt", "wb") as f:
        f.write(engine.serialize())

2. 看，错误就是这个......

Checking 0/3...
Checking 1/3...
Checking 2/3...
[TensorRT] WARNING: onnx2trt_utils.cpp:320: Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32.
mini_code.py:52: DeprecationWarning: Use build_serialized_network instead.
  engine = builder.build_engine(network, config)
[TensorRT] WARNING: Convolution + generic activation fusion is disable due to incompatible driver or nvrtc
[TensorRT] WARNING: TensorRT was linked against cuBLAS/cuBLAS LT 11.4.2 but loaded cuBLAS/cuBLAS LT 11.2.1
[TensorRT] WARNING: Detected invalid timing cache, setup a local cache instead
[TensorRT] ERROR: 4: [pluginV2Builder.cpp::makeRunner::680] Error Code 4: Internal Error (Internal error: plugin node ScatterND_12 requires 36 bytes of scratch space, but only 0 is available. Try increasing the workspace size with IBuilderConfig::setMaxWorkspaceSize().
)
Traceback (most recent call last):
  File "mini_code.py", line 55, in <module>
    f.write(engine.serialize())
AttributeError: 'NoneType' object has no attribute 'serialize'

预期结果

Checking 0/3...
Checking 1/3...
Checking 2/3...
mini_code.py:53: DeprecationWarning: Use build_serialized_network instead.
  engine = builder.build_engine(network, config)
[TensorRT] WARNING: Convolution + generic activation fusion is disable due to incompatible driver or nvrtc
[TensorRT] WARNING: TensorRT was linked against cuBLAS/cuBLAS LT 11.4.2 but loaded cuBLAS/cuBLAS LT 11.2.1
[TensorRT] WARNING: Detected invalid timing cache, setup a local cache instead
[TensorRT] WARNING: TensorRT was linked against cuBLAS/cuBLAS LT 11.4.2 but loaded cuBLAS/cuBLAS LT 11.2.1

截图

如果使用的是MLab HomePod，请填写

• 宿主机 cpu/ram/cuda设备： [比如. intel i5-9300H/8GB/GTX1650]
• 宿主机操作系统/内核版本/GPU驱动： [比如. ubuntu 20.04/5.4.0-77-generic/460.80 ]
• MLab HomePod版本 [比如： 2.0-pro]

上下文备注

# init module
class MyModule(torch.nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        ...

    def forward(self, x):
        y = x[:1] + 1
        return y

这种情况下不会报错
暂时猜测当索引切片操作在等号左边时tensorrt的onnxParser就会报错？

[MHGL]

please refer to NVIDIA/TensorRT#1363 (comment)

[gemfield]

添加 config.max_workspace_size = 1 << 30