当融合的卷积、激活和残差加法读取分步连接视图时,会产生错误的 FP16 结果

import sys import numpy as np import onnxruntime as ort import tensorrt as trt import torch import torch.nn as nn LEVEL = int(sys.argv[1]) if len(sys.argv) > 1 else 3 C = 256

class Act(nn.Module): def forward(self, x): return x * (0.5 * torch.tanh(0.5 * x) + 0.5)

class Repro(nn.Module): def init(self): super().init() self.cv1 = nn.Conv2d(C, 2 * C, 1) self.b1 = nn.Conv2d(C, C // 2, 3, padding=1) self.b2 = nn.Conv2d(C // 2, C, 3, padding=1) self.cv2 = nn.Conv2d(3 * C, C, 1) self.post = nn.Sequential(nn.Conv2d(C, C, 1), Act(), nn.Conv2d(C, C, 1)) self.act = Act()

def forward(self, x):
    a, b = self.cv1(x).split((C, C), 1)
    o = self.act(self.b2(self.act(self.b1(b)))) + b  # b feeds both the add and the concat below
    o = self.post(o)  # keeps b alive so it is still needed at the concat
    return self.cv2(torch.cat((a, b, o), 1))

torch.manual_seed(0) model = Repro().eval().cuda().half() x = (torch.randn(1, C, 20, 20) * 3).cuda().half() torch.onnx.export(model, x, "repro.onnx", input_names=["i"], output_names=["o"], opset_version=18) logger = trt.Logger(trt.Logger.ERROR) builder = trt.Builder(logger) cfg = builder.create_builder_config() cfg.builder_optimization_level = LEVEL cfg.profiling_verbosity = trt.ProfilingVerbosity.DETAILED net = builder.create_network(0) assert trt.OnnxParser(net, logger).parse_from_file("repro.onnx") eng = trt.Runtime(logger).deserialize_cuda_engine(builder.build_serialized_network(net, cfg)) ctx = eng.create_execution_context() bufs = [] for i in range(eng.num_io_tensors): nm = eng.get_tensor_name(i) dt = trt.nptype(eng.get_tensor_dtype(nm)) t = torch.zeros(tuple(ctx.get_tensor_shape(nm)), dtype=getattr(torch, np.dtype(dt).name), device="cuda") ctx.set_tensor_address(nm, t.data_ptr())