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

文章目录

1 数据预处理1.1 下载数据集1.2 制作bs_dataset1.3 制作bs_dataloader 2 搭建网络模型3 训练和测试模型4 总结

【本文完整代码：/cqfdch/BelgiumTSC-pytorch】

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

1 数据预处理

1.1 下载数据集

能拿来训练的数据集=数据集+标签文件

方法一：可以去官网https://btsd.ethz.ch/shareddata/下载数据集，然后进行自行处理

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

/download/Hankerchen/13073778

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

1.2 制作bs_dataset

根据pytorch标准的制作数据集的三段式

definit(self）:

deflen(self):

defgetitem(self, idx):

import torchimport osimport pandas as pdfrom torch.utils.data import Datasetimport numpy as npfrom PIL import Imageclass 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 appliedon a sample."""self.root_dir = root_dirself.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 = transformdef __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

参考链接/tomlawrenceuk/GTSRB-Dataloader

通过torch.utils.data.DataLoader（）将bs_dataset导入bs_loader。

import bs_dataset as datasetimport torchvision.transforms as transformsimport torchdef get_train_valid_loader(data_dir,batch_size,num_workers=0,):# Create Transformstransform = pose([transforms.Resize((32, 32)),transforms.ToTensor(),transforms.Normalize((0.3403, 0.3121, 0.3214),(0.2724, 0.2608, 0.2669))])# Create Datasetstrainset = dataset.BelgiumTSC(root_dir=data_dir, train=True, transform=transform)testset = dataset.BelgiumTSC(root_dir=data_dir, train=False, transform=transform)# Load Datasetstrainloader = 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

/cqfdch/BelgiumTSC-pytorch

3 训练和测试模型

参考链接

/qq_37541097/article/details/104710784

将前面制作好的bs_loader进行训练，每个epoch测试一次。

from __future__ import print_functionimport torchimport torch.nn as nnimport torch.optim as optimfrom tqdm import tqdmimport bs_loader# from model import AlexNetfrom model import Modelprint(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 = 30save_path = './model.pth'best_acc = 0.0for epoch in range(epoch):# trainnet.train()running_loss = 0.0total = 0correct = 0show_step = 32for step, data in enumerate(tqdm(train_loader),0):images, labels = datadevice = 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 statisticsrunning_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()# validatenet.eval()val_loss = 0.0total = 0correct = 0acc = 0.0 # accumulate accurate number / epochwith torch.no_grad():for _,val_data in enumerate(validate_loader,0):val_images, val_labels = val_datadevice = 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 总结

完整代码

/cqfdch/BelgiumTSC-pytorch

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