小技巧（6）：进行BelgiumTSC交通标志数据集识别（定义自己的数据集）

小技巧（5）：将TT100K数据集转成VOC格式，并且用Python脚本选出45类超过100张的图片和XML

【本文完整代码：https://github.com/cqfdch/BelgiumTSC-pytorch】

图像分类是计算机视觉的基础，pytorch做图像分类都是模块化的操作，主要包括数据预处理，搭建网络模型，训练和测试模型。

1 数据预处理

1.1 下载数据集

能拿来训练的数据集=数据集+标签文件
方法一：可以去官网https://btsd.ethz.ch/shareddata/下载数据集，然后进行自行处理

方法二：（1积分，看着支持就行）通过csdn资源下载（含数据集+标签文件）
https://download.csdn.net/download/Hankerchen/13073778

方法三：我的百度网盘（含数据集+标签文件）链接: https://pan.baidu.com/s/1JYWEFYFJCSRsVPmBfkauPQ 密码: wqtv

1.2 制作bs_dataset

根据pytorch标准的制作数据集的三段式
def init(self）:
def len(self):
def getitem(self, idx):

import torch
import os
import pandas as pd
from torch.utils.data import Dataset
import numpy as np
from PIL import Image


class BelgiumTSC(Dataset):
    base_folder = 'BelgiumTSC'

    def __init__(self, root_dir, train=False, transform=None):
        """
        Args:
            train (bool): Load trainingset or test set.
            root_dir (string): Directory containing GTSRB folder.
            transform (callable, optional): Optional transform to be applied
                on a sample.
        """
        self.root_dir = root_dir

        self.sub_directory = 'Training' if train else 'Testing'
        self.csv_file_name = 'train_data.csv' if train else 'test_data.csv'

        csv_file_path = os.path.join(
            root_dir, self.base_folder, self.sub_directory, self.csv_file_name)

        self.csv_data = pd.read_csv(csv_file_path)

        self.transform = transform

    def __len__(self):
        return len(self.csv_data)

    def __getitem__(self, idx):
        img_path = os.path.join(self.root_dir, self.base_folder, self.sub_directory,
                                self.csv_data.iloc[idx, 0])
        img = Image.open(img_path)

        classId = self.csv_data.iloc[idx, 1]

        if self.transform is not None:
            img = self.transform(img)

        return img, classId

1.3 制作bs_dataloader

参考链接https://github.com/tomlawrenceuk/GTSRB-Dataloader
通过torch.utils.data.DataLoader（）将bs_dataset导入bs_loader。

import bs_dataset as dataset
import torchvision.transforms as transforms
import torch

def get_train_valid_loader(data_dir,
                           batch_size,
                           num_workers=0,
                           ):
    # Create Transforms
    transform = transforms.Compose([
        transforms.Resize((32, 32)),
        transforms.ToTensor(),
        transforms.Normalize((0.3403, 0.3121, 0.3214),
                             (0.2724, 0.2608, 0.2669))
    ])

    # Create Datasets
    trainset = dataset.BelgiumTSC(
        root_dir=data_dir, train=True,  transform=transform)
    testset = dataset.BelgiumTSC(
        root_dir=data_dir, train=False,  transform=transform)

    # Load Datasets
    trainloader = torch.utils.data.DataLoader(
        trainset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
    testloader = torch.utils.data.DataLoader(
        testset, batch_size=batch_size, shuffle=False, num_workers=num_workers)
    return trainloader, testloader

2 搭建网络模型

网络模型有很多，AlexNet，VGGNet，ResNet等，注意要数据图片尺寸。
具体的代码可以看我的github
https://github.com/cqfdch/BelgiumTSC-pytorch

3 训练和测试模型

参考链接
https://blog.csdn.net/qq_37541097/article/details/104710784
将前面制作好的bs_loader进行训练，每个epoch测试一次。

from __future__ import print_function
import torch
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm
import bs_loader
# from model import AlexNet
from model import Model

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

train_loader, validate_loader = bs_loader.get_train_valid_loader('D:\\Networks\\cnn-ga-master\\data', batch_size=32, num_workers=0)

net = Model()
net.to(device)
loss_function = nn.CrossEntropyLoss()
# pata = list(net.parameters())
optimizer = optim.Adam(net.parameters(), lr=0.0002)
epoch = 30
save_path = './model.pth'
best_acc = 0.0
for epoch in range(epoch):
    # train
    net.train()
    running_loss = 0.0
    total = 0
    correct = 0
    show_step = 32
    for step, data in enumerate(tqdm(train_loader),0):
        images, labels = data
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = net(images)
        loss = loss_function(outputs, labels)
        loss.backward()
        optimizer.step()

        # print statistics
        running_loss += loss.item()
        _, predicted = torch.max(outputs.detach(), 1)
        total += labels.size(0)
        correct += predicted.eq(labels.data).sum().item()
        # print train process
        # rate = (step+1)/len(train_loader)
        # a = "*" * int(rate * 50)
        # b = "." * int((1 - rate) * 50)
        # print("\rtrain loss: {:^3.0f}%[{}->{}]{:.4f}".format(int(rate*100), a, b, loss), end="")
        # print('Train-Epoch:%3d, %3d / %3d ,Loss: %.3f, Acc:%.3f'% (epoch+1, step+1, len(train_loader),running_loss/total, (correct/total)))
        if step % show_step == 0:
            print("Epoch [{}][{}/{}]:Loss:{:.3f},Acc:{:.3f}".format(epoch+1, step+1, len(train_loader),running_loss/total, (correct/total)))
    print()

    # validate
    net.eval()
    val_loss = 0.0
    total = 0
    correct = 0
    acc = 0.0  # accumulate accurate number / epoch
    with torch.no_grad():
        for _,val_data in enumerate(validate_loader,0):
            val_images, val_labels = val_data
            device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
            val_images, val_labels = val_images.to(device), val_labels.to(device)
            outputs = net(val_images)
            loss = loss_function(outputs, val_labels)
            val_loss += loss.item()
            _, predicted = torch.max(outputs.detach(), 1)
            total += val_labels.size(0)
            correct += predicted.eq(val_labels.data).sum().item()
        if correct / total > best_acc:
            best_acc = correct / total
            # print('*'*100, self.best_acc)
            torch.save(net.state_dict(), save_path)
        print('Validate-Loss:%.3f, Acc:%.3f' % (val_loss / total, correct / total))


print('Finished Training')

4 总结

完整代码
https://github.com/cqfdch/BelgiumTSC-pytorch

觉得有帮助的可以给个赞，给个星！