加载pt模型输入增加permute层并保存模型

当你需要在PyTorch中加载预训练模型（例如`pt`模型），并且想要添加一个`Permute`层以便改变数据的维度顺序或通道顺序，你可以按照以下步骤操作：

1. 首先，确保你已经安装了必要的库，如`torchvision`如果模型来自` torchvision.models`。

import torch
from torch import nn
from torchvision.models import resnet18

# 加载预训练的resnet18模型
model = resnet18(pretrained=True)

2. 然后，找到你希望添加`Permute`层的位置。通常是在某个卷积层或全连接层之后。假设你想在最后一个卷积层之后添加`Permute`，你可以这样做：

# 查找模型的最后一层
last_layer = list(model.children())[-2]  # 对于ResNet，这是Conv2d层

# 创建一个新的Permute层，并设置你要交换的轴
permute_layer = nn.Permute(dims=(1, 2, 0))  # 这里通常是将(通道, 高度, 宽度)改为(宽度, 高度, 通道)

# 将原层替换为包含Permute的新序列
new_layer_sequence = [last_layer, permute_layer]
model.layers = new_layer_sequence  # 可能需要自定义属性layers，这里假设模型有这个属性

3. 最后，保存修改后的模型。如果你的模型原本是以`state_dict`形式存储的，可以这样做：

# 获取当前模型的状态字典
state_dict = model.state_dict()

# 更新状态字典，因为模型结构已变，我们需要移除原有模型中不存在的Permute层权重
for name, param in state_dict.items():
    if 'permute' in name:
        del state_dict[name]

# 将新的Permute层添加到state_dict
state_dict['new_permute_layer.weight'] = permute_layer.weight
state_dict['new_permute_layer.bias'] = permute_layer.bias

# 保存新的模型
new_model_name = 'permuted_resnet18.pth'
torch.save(state_dict, new_model_name)