第P9周：YOLOv5-Backbone模块实现-开发者社区

🍨 本文为🔗365天深度学习训练营中的学习记录博客
🍖 原作者：K同学啊

前言

实验环境

python 3.9.23 pytorch 2.5.1 pytorch-cuda 11.8 pytorch-mutex 1.0 torchaudio 2.5.1 torchinfo 1.8.0 torchvision 0.20.1 Visual Studio Code 1.104.2 (user setup)

代码实现

所用模块导入

importpathlibimportwarningsfromdatetimeimportdatetimeimporttorchimporttorch.nnasnnimporttorch.nn.functionalasFfromtorch.utils.dataimportDataLoader,random_splitfromtorchvisionimporttransforms,datasets# 可视化importmatplotlib.pyplotasplt

YOLOv5 Backbone 简介

YOLOv5 的主干网络（Backbone）负责从原始图像中提取多尺度特征。其核心组件包括：
- Conv 模块：标准卷积 + BatchNorm + SiLU 激活，这部分就像一个精细的过滤系统，它能够有效地识别并强调图像中的重要特征，同时抑制不相关的信息。标准卷积层负责捕捉输入数据的空间结构信息；批归一化有助于加速训练过程；SiLU激活函数则通过其平滑非线性的特性增强模型的表达能力。
- C3 模块：基于 CSP（Cross Stage Partial）结构，增强特征表达能力，这部分将前一层的输出与当前层的输出结合，有助于缓解深层网络中的梯度消失问题，使得网络能够更深更高效地学习复杂的模式。
- SPPF（Spatial Pyramid Pooling - Fast）：通过不同尺度的最大池化融合上下文信息，提升感受野，好比我们用不同的视角或缩放比例来观察一幅画，有时需要近距离细致观察局部细节，有时则需远距离把握整体构图。SPPF模块所做的就是让模型能够以多种尺度“观看”输入图像，确保无论目标物体大小如何，都能准确识别。

代码实现

设置gpu以及忽略警告

# 忽略警告信息warnings.filterwarnings("ignore")# 设置设备：优先使用 GPU（CUDA），否则使用 CPUdevice=torch.device("cuda"iftorch.cuda.is_available()else"cpu")print(f"使用设备:{device}")

数据准备

data_dir="./data/"data_path=pathlib.Path(data_dir)# 获取类别名称（从子文件夹名提取）class_names=sorted([item.nameforitemindata_path.iterdir()ifitem.is_dir()])print(f"检测到{len(class_names)}个类别:{class_names}")# 定义训练集变换train_transforms=transforms.Compose([transforms.Resize([224,224]),# 统一调整图像尺寸为 224x224transforms.ToTensor(),# 转为 Tensor，并归一化像素值到 [0, 1]transforms.Normalize(# 标准化：使用 ImageNet 的均值和标准差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])])# 使用 ImageFolder 加载全部数据（统一用 train_transform，后续再划分）# 注意：ImageFolder 要求目录结构为 root/class_name/*.jpgtotal_data=datasets.ImageFolder(root=data_dir,transform=train_transforms)# 划分训练集和测试集（8:2）total_size=len(total_data)train_size=int(0.8*total_size)test_size=total_size-train_size train_dataset,test_dataset=random_split(total_data,[train_size,test_size])# 为测试集重新应用无增强的 transform# 因为 random_split 不会改变 transform，需手动替换test_dataset.dataset.transform=test_transform# 创建 DataLoaderbatch_size=4train_loader=DataLoader(train_dataset,batch_size=batch_size,shuffle=True,num_workers=0)# Windows 下 num_workers>0 可能出错，设为0test_loader=DataLoader(test_dataset,batch_size=batch_size,shuffle=False,num_workers=0)# 打印一个 batch 的数据形状以验证forX,yintest_loader:print(f"输入张量形状 [Batch, Channel, Height, Width]:{X.shape}")print(f"标签形状:{y.shape}, 数据类型:{y.dtype}")break

模型构建

# 定义 YOLOv5 主干网络（用于分类）defautopad(k,p=None):"""自动计算 padding，使输出尺寸等于输入尺寸（same padding）"""ifpisNone:p=k//2ifisinstance(k,int)else[x//2forxink]returnpclassConv(nn.Module):"""标准卷积模块：Conv -> BatchNorm -> SiLU"""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()ifactisTrueelse(actifisinstance(act,nn.Module)elsenn.Identity())defforward(self,x):returnself.act(self.bn(self.conv(x)))classBottleneck(nn.Module):"""标准瓶颈层（带 shortcut）"""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=shortcutandc1==c2defforward(self,x):returnx+self.cv2(self.cv1(x))ifself.addelseself.cv2(self.cv1(x))classC3(nn.Module):"""YOLOv5 特色模块"""def__init__(self,c1,c2,n=1,shortcut=True,g=1,e=0.5):super().__init__()c_=int(c2*e)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_inrange(n)))defforward(self,x):returnself.cv3(torch.cat((self.m(self.cv1(x)),self.cv2(x)),dim=1))classSPPF(nn.Module):"""快速空间金字塔池化（SPPF）"""def__init__(self,c1,c2,k=5):super().__init__()c_=c1//2self.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)defforward(self,x):x=self.cv1(x)y1=self.m(x)y2=self.m(y1)y3=self.m(y2)returnself.cv2(torch.cat([x,y1,y2,y3],1))classYOLOv5_backbone(nn.Module):"""YOLOv5 主干网络 + 分类头"""def__init__(self,num_classes=4):super().__init__()# YOLOv5 Backbone 结构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=num_classes))defforward(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)returnx

模型实例化以及打印模型结构

# 实例化模型并移至设备model=YOLOv5_backbone(num_classes=len(class_names)).to(device)print("模型已加载到设备:",device)# 打印模型结构fromtorchsummaryimportsummary summary(model,(3,224,224))

构建训练函数和测试函数

# 训练与测试函数deftrain_one_epoch(dataloader,model,loss_fn,optimizer):"""单轮训练"""model.train()size=len(dataloader.dataset)num_batches=len(dataloader)total_loss,correct=0,0forX,yindataloader:X,y=X.to(device),y.to(device)pred=model(X)loss=loss_fn(pred,y)optimizer.zero_grad()loss.backward()optimizer.step()total_loss+=loss.item()correct+=(pred.argmax(1)==y).type(torch.float).sum().item()avg_acc=correct/size avg_loss=total_loss/num_batchesreturnavg_acc,avg_lossdefevaluate(dataloader,model,loss_fn):"""评估模型"""model.eval()size=len(dataloader.dataset)num_batches=len(dataloader)total_loss,correct=0,0withtorch.no_grad():forX,yindataloader:X,y=X.to(device),y.to(device)pred=model(X)total_loss+=loss_fn(pred,y).item()correct+=(pred.argmax(1)==y).type(torch.float).sum().item()avg_acc=correct/size avg_loss=total_loss/num_batchesreturnavg_acc,avg_loss

正式训练

# 开始训练# 超参数设置epochs=60learning_rate=1e-4# 损失函数与优化器loss_fn=nn.CrossEntropyLoss()optimizer=torch.optim.Adam(model.parameters(),lr=learning_rate)# 记录训练过程train_acc_list=[]train_loss_list=[]test_acc_list=[]test_loss_list=[]best_acc=0.0best_model_state=Noneprint("\n开始训练...\n")forepochinrange(epochs):# 训练train_acc,train_loss=train_one_epoch(train_loader,model,loss_fn,optimizer)# 测试test_acc,test_loss=evaluate(test_loader,model,loss_fn)# 保存最佳模型iftest_acc>best_acc:best_acc=test_acc best_model_state=model.state_dict()# 保存状态字典而非整个模型（更安全）# 记录train_acc_list.append(train_acc)train_loss_list.append(train_loss)test_acc_list.append(test_acc)test_loss_list.append(test_loss)# 打印日志current_lr=optimizer.param_groups[0]['lr']print(f"Epoch{epoch+1:2d}/{epochs}| "f"Train Acc:{train_acc*100:.1f}% | Train Loss:{train_loss:.3f}| "f"Test Acc:{test_acc*100:.1f}% | Test Loss:{test_loss:.3f}| "f"LR:{current_lr:.2e}")# 保存最佳模型best_model_path="./best_model.pth"torch.save(best_model_state,best_model_path)print(f"\n训练完成！最佳测试准确率:{best_acc*100:.2f}%，模型已保存至{best_model_path}")

结果可视化

# 可视化训练过程# 设置中文字体（避免中文乱码）plt.rcParams['font.sans-serif']=['SimHei']plt.rcParams['axes.unicode_minus']=Falseplt.rcParams['figure.dpi']=100# 获取当前时间用于图表标题current_time=datetime.now().strftime("%Y-%m-%d %H:%M:%S")# 绘制 Accuracy 和 Loss 曲线plt.figure(figsize=(12,4))# 准确率曲线plt.subplot(1,2,1)plt.plot(train_acc_list,label='训练准确率')plt.plot(test_acc_list,label='测试准确率')plt.xlabel('Epoch')plt.ylabel('Accuracy')plt.title(f'训练与测试准确率\n{current_time}')plt.legend()plt.grid(True)# 损失曲线plt.subplot(1,2,2)plt.plot(train_loss_list,label='训练损失')plt.plot(test_loss_list,label='测试损失')plt.xlabel('Epoch')plt.ylabel('Loss')plt.title(f'训练与测试损失\n{current_time}')plt.legend()plt.grid(True)plt.show()

模型评估

# 最终测试（加载最佳模型）print("\n正在加载最佳模型进行最终测试...")final_model=YOLOv5_backbone(num_classes=len(class_names)).to(device)final_model.load_state_dict(torch.load(best_model_path,map_location=device))final_acc,final_loss=evaluate(test_loader,final_model,loss_fn)print(f"最终测试结果 | 准确率:{final_acc*100:.2f}% | 损失:{final_loss:.3f}")

代码运行截图

学习总结

YOLOv5 的主干网络（Backbone）通过 Conv、C3 和 SPPF 三大核心模块高效提取多尺度特征：Conv 作为基础单元完成卷积、归一化与激活；C3 借鉴 CSP 思想，将输入分路处理，一路经 Bottleneck 残差块深度变换，另一路直接 bypass，再拼接融合，既保留原始信息又增强表达能力，有效缓解梯度消失；SPPF 则通过级联最大池化快速扩大感受野，捕获多尺度上下文。当然这次最大的收获是我在代码方面得到质的提升，过去我习惯直接复制和修改现成代码，而现在，我开始主动去思考代码的构建，拆解代码逻辑，最终基本能够完成编写。