这里先给出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')
制作不易，如有帮助请点赞一下哦！
