加载torchscript模型并在输入和输出添加permute层

加载预训练的 TorchScript 模型时，通常模型是以特定维度（如 NCHW 或 NHWC）设计的。如果需要在输入或输出时改变数据的维度，例如为了适应不同的硬件或API要求，你可以选择在加载模型后，在模型的前向传播路径上添加 `nn.functional.permute` 层。

举个例子，假设你有一个保存为 TorchScript 的模型 `model_traced`，其期望输入为 NCHW 形式，而实际应用需要 NHWC。在这种情况下，可以这样做：

import torch
from torch.jit import load

# 加载已转换为TorchScript的模型
model_traced = load('path_to_your_traced_model.pt')

class PermuteWrapper(torch.nn.Module):
    def __init__(self, model_traced):
        super().__init__()
        self.model = model_traced
        
    def forward(self, nhwc_input):
        # 添加 permute 层进行输入维度变换
        nchw_input = nhwc_input.permute(0, 3, 1, 2)
        output = self.model(nchw_input)
        
        # 可能需要再次 permute 输出，取决于输出预期的维度
        output_nhwc = output.permute(0, 2, 3, 1)
        return output_nhwc

# 创建并实例化新的模型，集成permute操作
model_with_permute = PermuteWrapper(model_traced)

# 现在你可以直接使用 model_with_permute 来处理 NHWC 输入
nhwc_data = ...  # 你的NHWC格式数据
output = model_with_permute(nhwc_data)

通过这种方式，你在不影响原始模型结构的情况下，仅在前向传播路径中添加了额外的维度变换步骤。