如何更改一个训练好的网络的部分架构+重新训练部分参数,冻结不需要的参数+配合学习器更改(以Deeplabv3+为例)

这里先给出deeplav3+的架构(只给出主体部分):

class DeepLabV3Plus(BaseNet):
    def __init__(self, backbone, nclass):
        super(DeepLabV3Plus, self).__init__(backbone)

        low_level_channels = self.backbone.channels[0]
        high_level_channels = self.backbone.channels[-1]

        self.head = ASPPModule(high_level_channels, (12, 24, 36))

        self.reduce = nn.Sequential(nn.Conv2d(low_level_channels, 48, 1, bias=False),
                                    nn.BatchNorm2d(48),
                                    nn.ReLU(True))

        self.fuse = nn.Sequential(nn.Conv2d(high_level_channels // 8 + 48, 256, 3, padding=1, bias=False),
                                  nn.BatchNorm2d(256),
                                  nn.ReLU(True),

                                  nn.Conv2d(256, 256, 3, padding=1, bias=False),
                                  nn.BatchNorm2d(256),
                                  nn.ReLU(True),
                                  nn.Dropout(0.1, False))

        self.classifier = nn.Conv2d(256, nclass, 1, bias=True)

    def base_forward(self, x):
        h, w = x.shape[-2:]

        c1, _, _, c4 = self.backbone.base_forward(x)

        c4 = self.head(c4)
        c4 = F.interpolate(c4, size=c1.shape[-2:], mode="bilinear", align_corners=True)

        c1 = self.reduce(c1)

        out = torch.cat([c1, c4], dim=1)
        out = self.fuse(out)

        out = self.classifier(out)
        out = F.interpolate(out, size=(h, w), mode="bilinear", align_corners=True)

        return out

假设刚开始我们的模型的nclass设置为15,整个网络模型已经经过了一次训练,保存的模型参数为best.pth 现在我想改变这个网络的最后一层self.classifier,将它的输出改为2个通道(原来是15个通道)

1 首先我们加载预训练模型

model = DeepLabV3Plus(backbone, nclass)
model.load_state_dict(torch.load('best.pth'))

2 然后,我们需要更改那一层的分类器

in_features = model.classifier.in_channels
new_out_features = 15  # 新的类别数量
#构建新的分类层,其他参数可以自定义
model.classifier = nn.Conv2d(in_features, new_out_features, kernel_size=1, bias=True)

3 对参数的可学习性进行设置(冻结 or no)

# 先冻结所有层的参数
for param in model.parameters():
    param.requires_grad = False
# 然后将新的层的参数可学习性设置为True
model.classifier.requires_grad = True  

4 优化器要做出相应改动(如果本身就是选择grad为true的参数进行更新,那不需要更改)

# 定义优化器,只更新分类器层的参数
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001, momentum=0.9)

这样,得到的model便可以进行使用了。

我利用gpt写了一段训练的完整流程代码(pass 部分自行填充即可):

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.transforms import Compose, Resize, ToTensor
from torchvision.datasets import YourDataset  # 替换为你的数据集类
from tqdm import tqdm

# 定义数据集
transform = Compose([
    Resize((256, 256)),
    ToTensor()
])
dataset = YourDataset(root='path/to/your/dataset', transform=transform)  # 替换为你的数据集路径

# 定义数据加载器
batch_size = 32
data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# 定义模型
class DeepLabV3Plus(nn.Module):
    def __init__(self, backbone, nclass):
        super(DeepLabV3Plus, self).__init__()
        # 你的模型定义

    def forward(self, x):
        # 前向传播逻辑
        pass

backbone = None  # 你的骨干网络
nclass = 10  # 类别数量
model = DeepLabV3Plus(backbone, nclass)

# 加载预训练模型
model.load_state_dict(torch.load('aaa.pth'))

# 修改分类器层
in_features = model.classifier.in_channels
new_out_features = 10  # 新的类别数量
model.classifier = nn.Conv2d(in_features, new_out_features, kernel_size=1, bias=True)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# 训练过程
num_epochs = 10
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
model.train()

for epoch in range(num_epochs):
    running_loss = 0.0
    for images, labels in tqdm(data_loader, desc=f'Epoch {epoch + 1}/{num_epochs}'):
        images, labels = images.to(device), labels.to(device)

        # 梯度清零
        optimizer.zero_grad()

        # 前向传播
        outputs = model(images)

        # 计算损失
        loss = criterion(outputs, labels)

        # 反向传播
        loss.backward()

        # 参数更新
        optimizer.step()

        running_loss += loss.item() * images.size(0)

    epoch_loss = running_loss / len(dataset)
    print(f'Epoch {epoch + 1}/{num_epochs}, Loss: {epoch_loss:.4f}')

print('Finished Training')

制作不易,如有帮助请点赞一下哦!

  • 27
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值