一、将PyTorch模型保存为TorchScript格式
1）构造一个pytorch2TorchScript.py，示例代码如下：
import torch
import torch.nn as nn
import argparse
from networks.seg_modeling import model as ViT_seg
from networks.seg_modeling import CONFIGS as CONFIGS_ViT_seg
import warnings
warnings.filterwarnings("ignore")def get_model(args):"""加载并配置模型。根据输入的配置参数创建模型并加载预训练的权重文件。Args:args: 命令行输入的参数集合。Returns:解包后的模型实例。"""# 根据配置文件选择模型config_vit = CONFIGS_ViT_seg[args.vit_name]# 根据模型类型配置跳跃连接与分类数if 'R50' in args.vit_name:config_vit.n_classes = args.num_classesconfig_vit.n_skip = 3else:config_vit.n_classes = args.num_classesconfig_vit.n_skip = 0# 配置模型的输入图像patch大小config_vit.patches["size"] = (args.vit_patches_size, args.vit_patches_size)# 实例化模型，并将其转移到GPU上model = ViT_seg(config_vit, img_size=args.img_size, num_classes=config_vit.n_classes).cuda()model = nn.DataParallel(model) # 使用多个GPU进行并行计算# 加载预训练的模型权重model.load_state_dict(torch.load(args.model_path))# 返回解包后的模型实例return model.moduledef export_to_torchscript(model, example_input, output_path):"""将模型导出为TorchScript格式。Args:model: 要导出的模型实例。example_input: 示例输入张量，用于模型的跟踪（tracing）。output_path: 导出后的TorchScript模型保存路径。"""model.eval() # 设置模型为推理模式# 使用trace方法将模型导出为TorchScripttraced_model = torch.jit.trace(model, example_input)# 将导出的模型保存为.pt文件traced_model.save(output_path)print(f"模型已保存为TorchScript格式: {output_path}")def main(args):# 加载模型model_to_export = get_model(args)# 创建一个示例输入张量 (大小为[BZ, channel, args.img_size, args.img_size])example_input = torch.randn(1, 1, args.img_size, args.img_size).cuda() # 模型导出的保存路径output_path = "torchscript/model_scripted.pt"# 导出模型为TorchScript格式export_to_torchscript(model_to_export, example_input, output_path)if __name__ == "__m
