第 P9 周:YOLOv5-Backbone 模块实现
·
👉声明
- 本文为 365 天深度学习训练营 内部学习记录。
- 数据集为天气识别数据集,仅用于学习交流,请勿对外传播。
- 原作者:K同学啊
- 本篇为个人在 P9 关卡上的实践记录。
第 P9 周:YOLOv5-Backbone 模块实现
本次我将利用 YOLOv5 算法中的 Backbone 模块搭建网络。理论部分会在语雀和公众号(K 同学啊)中展开,本次内容除网络结构部分外其余与上周基本一致。
🏡 我的环境
- 语言环境:Python 3.8(我的本机实跑为 3.11,流程一致)
- 编译器:Jupyter Lab
- 数据集:天气识别数据集(
cloudy/rain/shine/sunrise) - 深度学习环境:PyTorch(教程参考
torch==1.12.1+cu113,torchvision==0.13.1+cu113)
一、前期准备
1) 设置 GPU
如果设备支持 GPU 就使用 GPU,否则使用 CPU。
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision
from torchvision import transforms, datasets
import os, PIL, pathlib, warnings
warnings.filterwarnings("ignore")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device
2) 导入数据
教程里使用了:
classeNames = [str(path).split("\\")[1] for path in data_paths]
这在 Linux/WSL 下可能有兼容问题。更稳妥的方式是使用 path.name:
import pathlib
# 兼容多种目录:优先 9-data / p9-data,兜底 ../p3/data
candidate_roots = [pathlib.Path("9-data"), pathlib.Path("p9-data"), pathlib.Path("../p3/data")]
data_dir = next((p for p in candidate_roots if p.exists() and len([x for x in p.iterdir() if x.is_dir()]) > 0), None)
assert data_dir is not None, "未找到有效数据目录(请检查 9-data / p9-data)"
data_paths = list(data_dir.glob("*"))
classNames = [p.name for p in data_paths if p.is_dir()]
classNames
3) 数据预处理与读取
train_transforms = transforms.Compose([
transforms.Resize([224, 224]),
transforms.RandomHorizontalFlip(p=0.5),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
test_transform = transforms.Compose([
transforms.Resize([224, 224]),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
total_data = datasets.ImageFolder("./9-data/", transform=train_transforms)
total_data
4) 划分数据集
train_size = int(0.8 * len(total_data))
test_size = len(total_data) - train_size
train_dataset, test_dataset = torch.utils.data.random_split(total_data, [train_size, test_size])
batch_size = 4
train_dl = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=1)
test_dl = torch.utils.data.DataLoader(test_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=1)
for X, y in test_dl:
print("Shape of X [N, C, H, W]: ", X.shape)
print("Shape of y: ", y.shape, y.dtype)
break
二、搭建 YOLOv5 Backbone 模块模型
📌 K 同学啊提示:可以尝试通过增加/调整 Backbone 内的 C3 与 Conv 组合来提高准确率。
1) 搭建模型(Backbone 结构)
核心思想:复现 YOLOv5 6.0 Backbone 的典型模块 Conv、Bottleneck、C3、SPPF,并接入分类头。
import torch.nn.functional as F
def autopad(k, p=None): # kernel, padding
if p is None:
p = k // 2 if isinstance(k, int) else [x // 2 for x in k]
return p
class Conv(nn.Module):
# Standard convolution
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):
super().__init__()
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
self.bn = nn.BatchNorm2d(c2)
self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())
def forward(self, x):
return self.act(self.bn(self.conv(x)))
class Bottleneck(nn.Module):
# Standard bottleneck
def __init__(self, c1, c2, shortcut=True, g=1, e=0.5):
super().__init__()
c_ = int(c2 * e)
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c_, c2, 3, 1, g=g)
self.add = shortcut and c1 == c2
def forward(self, x):
return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
class C3(nn.Module):
# CSP Bottleneck with 3 convolutions
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
super().__init__()
c_ = int(c2 * e) # hidden channels
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c1, c_, 1, 1)
self.cv3 = Conv(2 * c_, c2, 1)
self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
def forward(self, x):
return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), dim=1))
class SPPF(nn.Module):
# Spatial Pyramid Pooling - Fast
def __init__(self, c1, c2, k=5):
super().__init__()
c_ = c1 // 2
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c_ * 4, c2, 1, 1)
self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
def forward(self, x):
x = self.cv1(x)
with warnings.catch_warnings():
warnings.simplefilter("ignore")
y1 = self.m(x)
y2 = self.m(y1)
return self.cv2(torch.cat([x, y1, y2, self.m(y2)], 1))
class YOLOv5_backbone(nn.Module):
def __init__(self):
super(YOLOv5_backbone, self).__init__()
self.Conv_1 = Conv(3, 64, 3, 2, 2)
self.Conv_2 = Conv(64, 128, 3, 2)
self.C3_3 = C3(128, 128)
self.Conv_4 = Conv(128, 256, 3, 2)
self.C3_5 = C3(256, 256)
self.Conv_6 = Conv(256, 512, 3, 2)
self.C3_7 = C3(512, 512)
self.Conv_8 = Conv(512, 1024, 3, 2)
self.C3_9 = C3(1024, 1024)
self.SPPF = SPPF(1024, 1024, 5)
self.classifier = nn.Sequential(
nn.Linear(in_features=65536, out_features=100),
nn.ReLU(),
nn.Linear(in_features=100, out_features=4),
)
def forward(self, x):
x = self.Conv_1(x)
x = self.Conv_2(x)
x = self.C3_3(x)
x = self.Conv_4(x)
x = self.C3_5(x)
x = self.Conv_6(x)
x = self.C3_7(x)
x = self.Conv_8(x)
x = self.C3_9(x)
x = self.SPPF(x)
x = torch.flatten(x, start_dim=1)
x = self.classifier(x)
return x
device = "cuda" if torch.cuda.is_available() else "cpu"
model = YOLOv5_backbone().to(device)
参数量统计:
try:
import torchsummary as summary
summary.summary(model, (3, 224, 224))
except ModuleNotFoundError:
try:
from torchinfo import summary
summary(model, input_size=(1, 3, 224, 224), depth=3)
except ModuleNotFoundError:
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Total params: {total_params:,}")
print(f"Trainable params: {trainable_params:,}")
三、训练模型
1) 编写训练函数
import copy
# 训练循环
def train(dataloader, model, loss_fn, optimizer):
size = len(dataloader.dataset) # 训练集大小
num_batches = len(dataloader) # 批次数
train_loss, train_acc = 0, 0
for X, y in dataloader: # 获取图片及其标签
X, y = X.to(device), y.to(device)
# 计算预测误差
pred = model(X)
loss = loss_fn(pred, y)
# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
train_acc += (pred.argmax(1) == y).type(torch.float).sum().item()
train_loss += loss.item()
train_acc /= size
train_loss /= num_batches
return train_acc, train_loss
2) 编写测试函数
def test(dataloader, model, loss_fn):
size = len(dataloader.dataset) # 测试集大小
num_batches = len(dataloader) # 批次数
test_loss, test_acc = 0, 0
with torch.no_grad():
for imgs, target in dataloader:
imgs, target = imgs.to(device), target.to(device)
target_pred = model(imgs)
loss = loss_fn(target_pred, target)
test_loss += loss.item()
test_acc += (target_pred.argmax(1) == target).type(torch.float).sum().item()
test_acc /= size
test_loss /= num_batches
return test_acc, test_loss
3) 正式训练
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
loss_fn = nn.CrossEntropyLoss()
epochs = 60
train_loss, train_acc = [], []
test_loss, test_acc = [], []
best_acc = 0
for epoch in range(epochs):
model.train()
epoch_train_acc, epoch_train_loss = train(train_dl, model, loss_fn, optimizer)
model.eval()
epoch_test_acc, epoch_test_loss = test(test_dl, model, loss_fn)
train_acc.append(epoch_train_acc)
train_loss.append(epoch_train_loss)
test_acc.append(epoch_test_acc)
test_loss.append(epoch_test_loss)
if epoch_test_acc > best_acc:
best_acc = epoch_test_acc
best_model = copy.deepcopy(model)
lr = optimizer.state_dict()["param_groups"][0]["lr"]
template = ("Epoch:{:2d}, Train_acc:{:.1f}%, Train_loss:{:.3f}, "
"Test_acc:{:.1f}%, Test_loss:{:.3f}, Lr:{:.2E}")
print(template.format(epoch + 1, epoch_train_acc * 100, epoch_train_loss,
epoch_test_acc * 100, epoch_test_loss, lr))
PATH = "./best_model.pth"
torch.save(best_model.state_dict(), PATH)
print("Done")

📌 如果将优化器换成 SGD,可以重点观察:
- 收敛速度是否变慢;
- 验证集准确率是否波动更大;
- 是否需要调整学习率与动量才能稳定训练。
四、结果可视化
1) Loss 与 Accuracy 图
from datetime import datetime
current_time = datetime.now()
plt.figure(figsize=(12, 3))
plt.subplot(1, 2, 1)
plt.plot(train_acc, label="Training Accuracy")
plt.plot(test_acc, label="Test Accuracy")
plt.legend(loc="lower right")
plt.title("Training and Validation Accuracy")
plt.xlabel(current_time)
plt.subplot(1, 2, 2)
plt.plot(train_loss, label="Training Loss")
plt.plot(test_loss, label="Test Loss")
plt.legend(loc="upper right")
plt.title("Training and Validation Loss")
plt.show()

2) 模型评估
best_model.load_state_dict(torch.load(PATH, map_location=device))
epoch_test_acc, epoch_test_loss = test(test_dl, best_model, loss_fn)
五、本周总结
本周基于天气识别四分类任务,完成了 YOLOv5 Backbone 思想在图像分类场景下的迁移实践。整体流程包括:数据预处理与 8:2 划分、Backbone 模块复现(Conv / Bottleneck / C3 / SPPF)、分类头构建、训练与测试函数封装、60 轮训练、最佳模型保存及可视化评估。
实验结果说明:YOLOv5 主干网络在分类任务中同样具备较强的特征提取能力,训练过程可稳定收敛,具备继续优化和做对照实验的价值。
后续优化方向:
- 调整
C3重复次数与通道规模,观察准确率与参数量变化; - 对比
Adam与SGD在收敛速度和泛化性能上的差异; - 继续优化分类头结构,尝试更轻量的池化与全连接设计;
- 增加混淆矩阵、分类别指标等评估方式,提升实验分析完整性。
**核心结论:**本周已完成从“检测主干模块”到“分类任务落地”的关键迁移验证,达成 P9 目标,并为后续结构改进提供了可复现实验基础。
AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。
更多推荐



所有评论(0)