P8:Pytorch实战:YOLOv5-C3模块实现

  • 本文为365天深度学习训练营 中的学习记录博客
  • 原作者:K同学啊|接辅导、项目定制

1 引言

本文利用 YOLOv5 算法中的C3模块搭建网络,实现天气识别。

2 前期工作

前期工作中包含数据处理、划分数据集等相关操作,由于在前面的文章中都有较为详细的解释,故在此只贴出代码。

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torchvision import datasets, transforms

import os, PIL, pathlib,random

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

import pathlib

data_dir = r'/Volumes/Monty _Lee/Learn/DL/Pytorch/P8/data'
data_dir = pathlib.Path(data_dir)

data_paths = list(data_dir.glob('*'))
classNames = [str(path).split('/')[-1] for path in data_paths]
print(classNames)

train_dir = data_dir / 'train'
test_dir = data_dir / 'test'
print(train_dir)
print(test_dir)

from torchvision import datasets, transforms

train_transforms = transforms.Compose([
    transforms.Resize([224,224]),
    transforms.RandomHorizontalFlip(), # 随机翻转和旋转
    transforms.ToTensor(),
    transforms.Normalize(mean = [0.485, 0.456, 0.406], # 均值
                         std = [0.229, 0.224, 0.225]) # 方差
])

test_transforms = 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_dataset = datasets.ImageFolder(data_dir, transform = train_transforms)

total_dataset.class_to_idx

train_size = int(0.8 * len(total_dataset))
test_size = len(total_dataset) - train_size
train_data, test_data = torch.utils.data.random_split(total_dataset, [train_size, test_size])

batch_size = 32
train_dl = torch.utils.data.DataLoader(train_data,
                                       batch_size = batch_size,
                                       shuffle = True,
                                       num_workers = 1,
                                       drop_last = True)
test_dl = torch.utils.data.DataLoader(test_data,
                                      batch_size = batch_size,
                                      shuffle = True,
                                      num_workers = 1,
                                      drop_last = True)


3 搭建包含C3模块的模型

import torch.nn.functional as F
from module import *

def autopad(k, p=None): # kernel, padding 
    # Pad to 'same' 
    if p is None: 
        p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad 
        
    return p 
class Conv(nn.Module): # Standard convolution 
    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups 
        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): # ch_in, ch_out, shortcut, groups, expansion 
        super().__init__() 
        c_ = int(c2 * e) # hidden channels 
        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): # ch_in, ch_out, number, shortcut, groups, expansion 
        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) # act=FReLU(c2) 
        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 model_K(nn.Module): 
    def __init__(self): 
        super(model_K, self).__init__() # 卷积模块 
        self.Conv = Conv(3, 32, 3, 2) # C3模块1 
        self.C3_1 = C3(32, 64, 3, 2) # 全连接网络层,用于分类

        self.classifier = nn.Sequential( 
            nn.Linear(in_features=802816, out_features=100), nn.ReLU(), 
            nn.Linear(in_features=100, out_features=4) 
            ) 
        
    def forward(self, x): 
        x = self.Conv(x) 
        x = self.C3_1(x) 
        x = torch.flatten(x, start_dim=1) 
        x = self.classifier(x) 
        return x 
        
device = "cuda" if torch.cuda.is_available() else "cpu" 
        
print("Using {} device".format(device)) 
model = model_K().to(device) 
model

上述代码定义了一个用于图像分类的神经网络模型。下面对代码进行解释:

  • Conv 类代表一个标准的卷积层,包括卷积操作、批归一化和激活函数(默认为 SiLU 激活函数)。
  • Bottleneck 类实现了一个标准的瓶颈块,包括一个 1x1 卷积、一个 3x3 卷积和一个可选的快捷连接。
  • C3 类是 CSP(Cross Stage Partial)瓶颈块,是瓶颈块的改进版本。它包括两个 1x1 卷积、一系列的瓶颈块以及一个 1x1 卷积来将结果拼接起来。
  • model_K 类定义了整个模型的架构。它包括一个 Conv 层,然后是一个 C3 块。最后,有一个用于分类的全连接网络(FCN),包括两个线性层。

3.1 YOLOv5

YOLOv5 是一种常见的目标检测算法和架构,它基于之前的 YOLO(You Only Look Once)模型进行改进和扩展。YOLOv5 由 Ultralytics 开发,并在计算机视觉领域引起了广泛关注。

相比于前几个版本的 YOLO,YOLOv5 进行了几个关键的修改,包括:

  1. 架构:YOLOv5 使用了与 YOLOv4 不同的修改后的架构。它采用了 CSPDarknet53 主干网络,这是 Darknet 架构的修改版本,用于特征提取。

  2. 主干网络:CSPDarknet53 主干网络由卷积层组成,并利用交叉阶段部分连接来增强网络中不同阶段之间的信息流动。

  3. 目标检测头部:YOLOv5 使用了定制的目标检测头部,用于预测边界框、类别概率和目标性分数。它利用不同尺度和长宽比的锚框来处理不同大小的目标物体。

  4. 模型缩放:YOLOv5 引入了一个模型缩放系统,允许用户根据需求在速度和准确性之间进行权衡。模型可以按照不同尺寸进行缩放,例如小型、中型、大型和特大型,以适应不同的部署场景。

  5. 训练方法:YOLOv5 使用了一种名为“自动学习率调度”(Automated Learning Rate Scheduling,ALRS)的新型训练策略。ALRS 在训练过程中自动调整学习率,以实现更快的收敛和更好的性能。

YOLOv5 在各种目标检测基准测试中表现出色,具有实时目标检测能力和较高的准确性。它已被广泛应用于自动驾驶、监控系统以及图像和视频中的目标识别等各种应用领域。

3.2 C3模块

在YOLOv5中,C3模块是一种重要的组件,它是一种CSP(Cross Stage Partial)瓶颈块。下面是对C3模块的简要介绍:

C3模块是YOLOv5中用于特征提取的重复模块。它基于瓶颈结构,并包含三个卷积层。C3模块的设计目的是在不增加计算量的情况下提高特征表达能力。

C3模块的输入是一个特征张量,通常是来自主干网络的输出。接下来,输入特征通过两个并行的1x1卷积层进行处理。这两个1x1卷积层用于降低特征的通道数,并生成两个具有较小通道数的特征张量。

然后,这两个特征张量被用作输入,通过一系列堆叠的瓶颈块(Bottleneck)进行处理。每个瓶颈块由一次1x1卷积和一次3x3卷积组成,它们的目标是进一步提取和增强特征信息。

最后,经过一系列的瓶颈块之后,两个特征张量被拼接在一起,并通过一个1x1卷积层来融合它们。这样,C3模块的输出特征张量就得到了,可以继续用于后续的处理,例如目标检测头部的操作。

总的来说,C3模块在YOLOv5中起到了提取丰富特征并增强特征表达能力的作用。它通过使用瓶颈块和1x1卷积层的组合来实现这一目标,并在保持计算效率的同时提高了模型的性能。

4 训练模型

4.1 训练函数和测试函数

# 训练函数
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)

        # Compute prediction error
        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

# 测试函数
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 x, y in dataloader:
            x, y  = x.to(device), y.to(device)

            pred = model(x)
            loss = loss_fn(pred, y)

            test_loss += loss_fn(pred, y).item()
            test_acc += (pred.argmax(1) == y).type(torch.float).sum().item()

    test_acc /= size
    test_loss /= num_batches

    return test_acc, test_loss


4.2 正式训练

import copy
from tqdm import tqdm

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()


epoches = 10
loss_fn = nn.CrossEntropyLoss()
train_loss = []
train_acc = []
test_loss = []
test_acc = []

best_acc = 0.0

for epoch in tqdm(range(epoches)):
    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)

    if epoch_test_acc > best_acc:
        best_acc = epoch_test_acc
        best_model_wts = copy.deepcopy(model.state_dict())


    train_acc.append(epoch_train_acc)
    train_loss.append(epoch_train_loss)
    test_acc.append(epoch_test_acc)
    test_loss.append(epoch_test_loss)

    template = 'Epoch:{:2d}, Train_acc:{:.1f}%, Train_loss:{:.3f}, Test_acc:{:.1f}%, Test_loss:{:.3f}'
    tqdm.write(template.format(epoch+1, epoch_train_acc*100, epoch_train_loss, epoch_test_acc*100, epoch_test_loss))

PATH = r'E:\Learn\DL\Pytorch\P8\result\model.pth'
torch.save(model.state_dict(), PATH)

print('Done')

5 结果可视化

import matplotlib.pyplot as plt

epochs_range = range(epoches)
# print(epochs_range)

plt.figure(figsize = (12, 3))
plt.subplot(1, 2, 1)
plt.plot(epochs_range, train_acc, label = 'Training Accuracy')
plt.plot(epochs_range, test_acc, label = 'Test Accuracy')
plt.xlim((0,20))
plt.xticks(range(0,epoches+10,10))
plt.legend()
plt.title('Training and Validation Accuracy')

plt.subplot(1, 2, 2)
plt.plot(epochs_range, train_loss, label = 'Training Loss')
plt.plot(epochs_range, test_loss, label = 'Test Loss')
plt.xlim((0,20))
plt.xticks(range(0,epoches+10,10))
plt.legend()
plt.title('Training and Validation Loss')
plt.show()

P8:Pytorch实战:YOLOv5-C3模块实现_第1张图片

你可能感兴趣的:(Pytorch,pytorch,YOLO,人工智能)