使用 PyTorch、ONNX 和 TensorRT 将视觉 Transformer 预测速度提升 9 倍

U-NET、Swin UNETR 等视觉转换器在语义分割等计算机视觉任务中是最先进的。

• U-NET 是弗赖堡大学计算机科学系为生物医学图像分割开发的卷积神经网络。其基于完全卷积网络，并在结构上加以修改与扩展，使得它可以用更少的训练图像产生更精确的分割。在现代GPU上，分割一张512×512的图像需要的时间不到一秒。
• Swin UNETR

但此类模型需要花费大量时间才能做出预测。本文展示了如何将此类模型的预测速度加快 9 倍。这一改进为一些实时或近实时应用程序提供了一种解决思路。

肿瘤分割任务

为了更好说明问题这里设置了一个场景，使用 Swin UNETR 模型从胸部 CT 扫描图像（单通道灰度 3D 图像）中分割肺部肿瘤。这是一个例子：

上述图片来源：wiki.cancerimagingarchive.net/display/Pub…

• 左栏显示了 3D CT 扫描图像在轴向平面上的一些 2D 切片。两个新月形的黑色区域是肺部。
• 右栏显示肺部肿瘤的手动标注。

胸部 CT 扫描的大小通常为 512×512×300，大约需要 60 到 90 兆字节存储在磁盘中。它们不是小图像。

使用 PyTorch 训练 Swin UNETR 模型来分割肺部肿瘤。经过训练的模型大约需要 10 秒才能对胸部 CT 扫描进行预测。所以每张图像 10 秒是起点。

在讨论速度优化之前，先看看模型的输入和输出，以及它如何进行预测。

模型输入和输出图像

• 输入是胸部 CT 扫描的 3D numpy 数组。
• Swin UNETR 模型无法容纳整个图像（太大了）。解决方案是将图像切割成更小的块，称为感兴趣区域 ROI。在这里设置感兴趣区域的大小为 96×96×96。
• Swin UNETR 模型一次看到一个感兴趣区域，输出两个二元分割掩模，一个用于肿瘤类别，另一个用于背景类别。两个掩模都是感兴趣区域的大小，即 96×96×96。更准确地说，Swin UNETR 输出两个非标准化类别概率掩码。在后面的步骤中，这些未归一化的掩码通过 softmax 归一化为 0 到 1 之间的适当概率，然后通过 argmax-ed 转换为二进制掩码。
• 这些掩模根据相应感兴趣区域的切割方式进行合并，以提供两个全尺寸分割掩模——肿瘤掩模和背景掩模——每个掩模具有整个胸部CT扫描的尺寸。请注意，即使模型返回两个分割掩模，这里只对肿瘤掩模感兴趣，并且会忽略背景掩模。
• 输入和输出数组以及模型使用 32 位浮点。

滑动窗口推理

下面的伪代码实现了上述预测思想。

def sliding_window_inference(image:np.ndarray,model:nn.Module,batch_size=4):
    rois = split_image(image)
    batches = [rois[i:i + batch_size] for i in range(0,len(rois),batch_size)]
    predictions_for_rois = [model(batch) for batch in batches]
    lesion_mask,background_mask = merge_predictions(predictions_for_rois)
    return lesion_mask,background_mask

请注意，本文中的代码片段是伪代码，以保持简洁。同样的道理，那些实现显而易见的方法，比如 split_image，就留给你想象吧。

• sliding_window_inference 方法接受完整 CT 扫描图像和 PyTorch 模型。它还接受 batch_size 因为感兴趣的区域很小，并且 GPU 可以一次容纳多个区域进行预测。 batch_size 指定发送到 GPU 的感兴趣区域的数量。 sliding_window_inference 返回二元肿瘤分割和背景掩模。
• 该方法首先将整个图像分割为感兴趣区域，然后将它们分组为批次，每个组包含 batch_size 感兴趣的区域。在这里，为了代码简单起见，我假设感兴趣区域的数量可以除以 batch_size。
• 每个批次都会发送到模型一批预测。每个预测都是针对单个感兴趣区域。
• 最后，所有感兴趣区域的预测被合并以形成两个全尺寸的分割掩模。合并还包括 softmax 和 argmax。

用于对图像进行预测的片段

以下代码片段调用 sliding_window_inference 方法对加载到第一个 GPU cuda:0中的图像文件进行预测： PyTorch 张量：

devices = 'cuda:0'
pretrained_pth="model.pt"
model_dict = torch.load(pretrained_pth)["state_dict"]
model = SwinUNETR()
model.load_state_dict(model_dict,strict=False)
model = model.to(device)
image = torch.Tensor(np.load('image.npy')).to(device)
lesion_mask,background_mask = sliding_window_inference(image,model)

通过上述设置，接下来介绍一组策略来使模型预测更快。

策略1：在 16bit 浮点数中进行预测

默认情况下，经过训练的 PyTorch 模型使用 32bit 浮点。但通常 16bit 浮点精度足以提供非常相似的分割结果，只需使用一个 PyTorch API 即可轻松将 32bit 模型转换为 16bit 模型：

image = image.half()
model = model.half()
lesion_mask,background_mask = sliding_window_inference(image,model)

这种策略将预测时间从 10 秒 减少到 7.7 秒。

策略 2：将模型转换为 TensorRT

TensorRT 是 Nvidia 的一个工具库，旨在为深度学习模型提供快速推理。它通过将在许多硬件中运行的通用模型（例如 PyTorch 模型或 TensorFlow 模型）转换为仅在一种特定硬件（即运行模型转换的硬件）中运行的 TensorRT 模型来实现此目的。在转换过程中，TensorRT 还执行许多速度优化。

TensorRT 安装中的 trtexec 可执行文件执行转换。问题是，有时从 PyTorch 模型到 TensorRT 模型的转换会失败。具体的失败消息并不重要，通常会遇到自己的错误。

重要的是要知道有一个绕行路线。解决方法是首先将 PyTorch 模型转换为中间格式 ONNX，然后将 ONNX 模型转换为 TensorRT 模型。

ONNX 是一种开放格式，旨在表示机器学习模型。 ONNX 定义了一组通用运算符（机器学习和深度学习模型的构建块）和通用文件格式，使 AI 开发人员能够将模型与各种框架、工具、运行时和编译器结合使用。

好消息是，将 ONNX 模型转换为 TensorRT 模型的支持比将 PyTorch 模型转换为 TensorRT 模型更好。

将 PyTorch 模型转换为 ONNX 模型

以下代码片段将 PyTorch 模型转换为 ONNX 模型：

device = "cuda:0"
dummy_input = torch.randn((1,1,96,96,96)).float().to(device) # A Single ROI
torch.onnx.export(
    model,
    dummy_input,
    'swinunetr.onnx'
    export_params=True,
    opset_version=17,
    do_constant_folding=True,
    input_names=['modelInput'],
    output_names = ['modelOutput'],
    dynamic_axes={
        'modelInput':{0:'dynamic'},
    }
)

它首先为单个感兴趣区域创建随机输入。然后使用已安装的 onnx Python 包中的导出方法来执行转换。这个转换输出一个名为 swinunetr.onnx 的文件。参数 dynamic_axes  指定TensorRT模型应该支持输入的第 0 维（即批处理维度）的动态大小。

将 ONNX 模型转换为 TensorRT 模型

现在可以调用 trtexec 命令行工具将 ONNX 模型转换为 TensorRT 模型：

trtexec --onnx=swinunetr.onnx --saveEngine=swinunetr_1_8_16.plan --fp16 --verbose --minShapes=modelInput:1×1×96×96×96 --optShapes=modelInput:8×1×96×96×96 --maxShapes=modelInput:16×1×96×96×96 --workspace=10240

• onnx=swinunetr.onnx 命令行选项指定 onnx 模型的位置。
• saveEngine=swinunetr_1_8_16.plan 选项指定生成的 TensorRT 模型的文件名，称为 plan。
• fp16 选项要求转换后的模型以 16bit 浮点精度运行
• minShapes=modelInput:1×1×96×96×96 指定生成的 TensorRT 模型的最小输入大小。
• maxshapes=modelInput:16×1×96×96×96 指定生成的 TensorRT 模型的最大输入大小。由于在 PyTorch 到 ONNX 转换过程中，只允许第 0 维，即批量维度支持动态大小，这里在 minShapes 和 maxShapes 中，只有第一个数字可以改变。它们一起告诉 trtexec 工具输出一个模型，该模型可用于批量大小在 1 ~ 16 之间的输入。
• optShapes=modelInput:8×1×96×96×96 指定生成的 TensorRT 模型应以批量大小 8 运行最快。
• workspace=10240 选项为 trtexec 提供 10G 的 GPU 内存来处理模型转换。

trtexec 将运行 10 到 20 分钟，并输出生成的 TensorRT  plan文件。

使用 TensorRT 模型进行预测

以下代码片段加载 TensorRT 模型 plan 文件并使用改编自 stackoverflow 的 TrtModel：

engine = 'swinunetr_1_8_16.plan'
device = 'cuda:0'
model = TrtModel(engine,dtype=np.float32,max_batch_size=4,device=device)
lesion_mask,background_mask = sliding_window_inference(image,model)

在 trtexec 命令行中，指定了 fp16 选项，但在加载 plan 时，仍然需要指定 32bit 浮点。从 stackoverflow 获得的 TrtModel 需要进行一些小的调整。采用此战术，预测时间为 2.89秒！

策略 3：封装模型以返回一个掩模

Swin UNETR 模型以非归一化概率的形式返回两个分割掩模，一个用于肿瘤，一个用于背景，以非标准化概率的形式。 这两个掩码首先从 GPU 传输回 CPU。 然后在 CPU中，这些非归一化概率经过softmax-ed 处理为 0 到 1 之间的适当概率，最后经过 argmax-ed 生成二进制掩码。由于只使用肿瘤掩模，因此模型不需要返回背景掩模。 GPU 和 CPU 之间传输数据需要时间，softmax 等计算也需要时间。

要拥有仅返回单个掩码的模型，可以创建一个封装 SwinUNETR 模型的新类：

class SwinWrapper(nn.Module):
    def __init__(self,model):
        super().__init__()
        self.model = model
    def forward(self,rois):
        # rois is of shape Batch × Class × Width × Height × Depth
        out = self.model(rois)
        out1 = F.softmax(out,dim=1)
        out2 = out1[:,1:2,:,:,:] # B,1,W,H,D
        return out2 # B,1,W,H,D

下图是新模型的输入输出流程：

方法 forward 通过神经网络的前向传递推送一批输入的感兴趣区域来进行预测。在这个方法中：

• 首先在传入的输入感兴趣区域上调用原始模型，以获得两个分割类的预测。输出的形状为 Batch×2×Width×Height×Depth，因为在当前的肿瘤分割任务中，有两类：肿瘤和背景。结果存储在 out 变量中。
• 然后将 softmax 应用于两个未归一化的分割掩码，将它们转换为 0 到 1 之间的归一化概率。
• 然后只选择肿瘤类别，即类别 1，将相应的值返回给调用者。

所以，实际上，这个封装实现了两个优化：

• 只返回一个分段掩码，而不是两个。
• 将 softmax 运算从 CPU 移至到 GPU。

argmax 操作会怎么样？由于只返回一个分段掩码，因此不需要 argmax 了。相反，为了创建原始的二进制分割掩码，将使 tumour_segmentation_probability ≥ 0.5，其中 tumour_segmentation_probability 是 SwinWrapper 中方法 forward 的结果。

由于 SwinWrapper 是 PyTorch 模型，因此需要再次执行 PyTorch 到 ONNX、ONNX 到 TensorRT 的转换步骤。

将 SwinWapper 模型转换为 ONNX 模型时，唯一变化是 model 需要经过 SwinWapper 处理：

device = "cuda:0"
dummy_input = torch.randn((1,1,96,96,96)).float().to(device) # A Single ROI
torch.onnx.export(
    SwinWrpper(model),
    dummy_input,
    'swinunetr.onnx'
    export_params=True,
    opset_version=17,
    do_constant_folding=True,
    input_names=['modelInput'],
    output_names = ['modelOutput'],
    dynamic_axes={
        'modelInput':{0:'dynamic'},
    }
)

将 ONNX 模型转换为 TensorRT 计划的 trtexec 命令行保持不变。

这种策略将预测时间从 2.89 秒 减少到 2.42 秒。

策略 4：将感兴趣区域分配给多个 GPU

上述所有策略仅使用一个 GPU，但有时希望使用更昂贵的多 GPU 机器来提供更快的预测。这个想法是将相同的 TensorRT 模型加载到 n 个 GPU 中，并且在 slider_window_inference 中，进一步将一批 ROI 拆分为 n 个部分，并将每个部分分配到不同的 GPU。这样，SwinWrapper 网络的耗时前向传递可以针对不同部分同时运行。

需要将 sliding_window_inference 方法更改为以下 sliding_window_inference_multi_gpu：

def sliding_window_inference_multi_gpu(image,models,batch_size,executor:ThreadPoolExecutor):
    rois = split_image(image)
    batches = [rois[i:i+batch_size] for i in range(0,len(rois),batch_size)]
    predictions_for_rois = []
    for batch in batches:
        futures = []
        for gpu_id,batch_per_gpu in enumerate(split_batch(batch,models)):
            fu = executor.submit(models[gpu_id],batch_per_gpu)
            futures.append(fu)
        done,not_done = wait(futures,return_when=ALL_COMPLETED)
        predictions = [fu.result() for fu in features]
        predictions_for_rois.append(np.concatenate(predictions))
    lesion_mask = merge_predictions(predictions_for_rois)
    return lesion_mask

• 和前面一样，将感兴趣的区域分组为不同的批次。
• 根据给定的 GPU 数量将每个批次分成几部分。
• 对于每个 batch_per_gpu 部分，将一个任务提交到 ThreadPoolExecutor 中，任务对传入的部分执行模型推理。
• 方法 submit 返回一个 future 对象，表示任务完成时的结果。 submit 方法在任务完成之前立即返回至关重要，因此可以将其他任务发布到不同的线程而无需等待，从而实现并行性。
• 在内部 for 循环中提交所有任务后，等待所有未来对象完成。
• 任务完成后，从 future 中读取结果并合并结果。

下面是调用 sliding_window_inference_multi_gpu 的代码片段：

model0 = TrtModel(engine,dtype=np.float32,max_batch_size=4,device='cuda:0')
model1 = TrtModel(engine,dtype=np.float32,max_batch_size=4,device='cuda:1')
models = [model0,model1]
executor = ThreadPoolExecutor(max_workers=2)
lesion_mask = sliding_window_inference_multi_gpu(image,models,batch_size=8,executor=executor)

• 这里使用了两个 GPU，因此创建了两个 TensorRT 模型，每个模型进入不同的 GPU，cuda:0 和 cuda:1。
• 然后创建了一个带有两个线程的 ThreadPoolExecutor。
• 将模型和执行器传递到 sliding_window_inference_multi_gpu 方法中，类似于单个GPU的情况，以获得肿瘤类分割掩模。

这种策略将预测时间从 2.42 秒 减少到 1.38 秒！

上面介绍了四种优化策略，将 SwinUNETR 模型的预测速度提高了 9 倍。

是否会为了速度而牺牲预测精度？

注意这里的 “精度” 是指最终模型分割肿瘤的程度，并不是指浮点预测，例如 16bit、32bit精度。

为了回答这个问题，需要看看衡量细分模型性能的 DICE 指标。

DICE 分数计算为预测肿瘤与真实肿瘤之间重叠的比例。 DICE分数在 0 到 1之间； DICE 分数越大意味着模型预测越好：

• DICE 为 1 是完美的预测
• DICE 为 0 是完全错误的预测，或者根本没有预测。

下面来看一下通过上述四种策略预测图像的 DICE 分数：

可以看到，只有当在策略1中将 32bit PyTorch模型转为 16bit 模型时，DICE分数从 0.93 略有下降到 0.91。其他策略不会进一步降低 DICE 分数。这表明该策略可以实现更快的预测速度，而精度损失却很小。

总结

本文介绍了四种策略，通过使用 ONNX、TensorRT 和多线程等工具使视觉转换器以更快的速度进行预测。

译自：towardsdatascience.com/making-visi…