49、PyTorch模型部署与C++应用实践

PyTorch模型部署与C++应用实践

1. JIT模块的使用与模型跟踪

JIT（Just-In-Time）编译的模块在Python中的工作方式与普通模块类似，既可以用于训练，也可用于推理。在进行推理时，需要像传统模型一样设置，例如使用 torch.no_grad() 上下文。

对于算法相对简单的模型，如CycleGAN、分类模型和基于U-Net的分割模型，可以直接进行模型跟踪。而对于更复杂的模型，有一个实用的特性：可以在其他脚本化或跟踪的代码中使用脚本化或跟踪的函数，在构建和跟踪或脚本化模块时使用脚本化或跟踪的子模块。也可以通过调用 nn.Models 来跟踪函数，但需要将所有参数设置为不需要梯度，因为对于跟踪的模型，这些参数将是常量。

2. 可跟踪性的脚本化处理

在一些复杂模型中，如Fast R-CNN系列的检测模型或自然语言处理中使用的循环网络，包含控制流（如for循环）的部分需要进行脚本化处理。以U-Net模型中的 UNetUpBlock 类为例：

class UNetUpBlock(nn.Module):
    ...
    def center_crop(self, layer, target_size):
        _, _, layer_height, layer_width = layer.size()
        diff_y = (layer_height - target_size[0]) // 2
        diff_x = (layer_w