背景介绍
本文是自己 使用图片数据集训练深度神经网络实现图片分类 的一次学习记录。其中数据集是b站小土堆评论区的蚂蚁蜜蜂数据集，训练和验证数据集总共只有400张左右的图片。数据集非常小，其实并不适合进行深度学习的训练。我只是想通过这个数据集来学习自己写dataset类并训练模型的整个过程。模型搭建使用的是PyTorch框架，模型是自己随便搭的，很简单，只有卷积（CNN）、最大池化（MaxPooling）、线性层（Linear）。如果有uu也打算写代码练习，可以b站搜小土堆获取数据集哈！
写这篇文章，一来是对自己编写代码过程的一次总结反思，二来详细记录实现模型训练的整个过程，既加深自己的记忆，又方便以后回忆学习。如果，其中所记录的点有问题，还烦请朋友们不要吝啬文字，欢迎在评论区指出、讨论！如果这篇学习记录对您有所启发的话，还请点赞收藏啦，感谢！！！
dataset类的实现
如果想使用自己的数据集，通常需要自己实现数据集的类。实现思路：写一个数据集类，继承自torch.utils.data.Dataset；在类里重写重写__init__、__getitem__、__len__方法，可见下方代码
from torch.utils.data import Dataset
import torchvision.transforms as transforms
import os
import cv2
trans = transforms.Compose([
transforms.ToTensor(),
# transforms.CenterCrop(256)
])
# 数据集继承自Dataset，重写init、getitem、len方法
class AB(Dataset):
def __init__(self, root, label):
self.root = root
self.label = label
self.data_path = os.path.join(self.root, self.label)
self.data = os.listdir(self.data_path)
def __getitem__(self, idx):
img_path = os.path.join(self.data_path, self.data[idx])
img = cv2.imread(img_path)
# 猜测img.resize()将img当作纯numpy数据处理，cv2.resize()把img当作图像
img = cv2.resize(img, (256, 256), interpolation=cv2.INTER_AREA) #后续会出一篇文章介绍cv2.resize()的插值方式
img = trans(img) # 在我所使用的训练集下，transforms.CenterCrop(256)的训练结果要比img = cv2.resize(img, (256, 256), interpolation=cv2.INTER_AREA)的训练结果差很多
if self.label == "ants":
return img, 0 # 返回的img是tensor
else:
return img, 1 # 由于是2分类，这里直接使用了if-else，如果类别比较多，可以使用字典实现
def __len__(self):
return len(self.data)
if __name__ == "__main__":
root = r"../../hymenoptera_data/train"
ant_label = "ants"
bee_label = "bees"
ants = AB(root, ant_label)
bees = AB(root, bee_label)
img = bees[9]
print(img[0].shape)
# cv2.imshow("Demo1", img)
dataset = ants + bees # 直接使用 + 实现数据集的合并
print(len(ants))
print(len(bees))
print(len(dataset))模型搭建
模型是自己随便搭的，很简单，只有卷积（CNN）、最大池化（MaxPooling）、线性层（Linear）。模型输入的图片是[ 3, 256, 256 ]
import torch.nn as nn
class LeNet5(nn.Module):
def __init__(self):
super().__init__() # 超继承
# 使用 nn.Sequential 组合操作
self.model = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=8, kernel_size=4, padding=1, stride=2),
nn.Conv2d(in_channels=8, out_channels=16, kernel_size=4, padding=1, stride=2),
nn.MaxPool2d(kernel_size=2),
nn.Conv2d(in_channels=16, out_channels=16, kernel_size=4, padding=1, stride=2),
nn.MaxPool2d(kernel_size=2),
nn.Flatten(),
nn.Linear(in_features=16 * 8 * 8, out_features=64),
nn.Linear(in_features=64, out_features=2) # 最后的二分类
)
def forward(self, input):
input = self.model(input)
return input
if __name__ == "__main__":
model = LeNet5()
print(model)训练脚本
训练的代码编写，要包括：模型、数据集、loss函数、优化器、学习率lr、batchsize、epoch、device（cuda or cpu）、训练参数保存。这一份代码里，把以上所有都考虑到，一份至少完整的训练脚本就写出来了。
接下来准备专题学习loss函数和优化器，如果有感兴趣的uu可以留下个关注以防走丢
import os
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as transforms
from torch.utils.tensorboard import SummaryWriter
from model import *
from predata import *
device = "cuda" if torch.cuda.is_available() else "cpu"
res = "./res" # 训练结果
if not os.path.exists(res):
os.makedirs(res)
writer = SummaryWriter("./logs") # 使用tensorboard保存训练、验证误差，观察训练情况
train_root = r"../../hymenoptera_data/train"
test_root = r"../../hymenoptera_data/val"
ant_label = "ants"
bee_label = "bees"
ant_train = AB(train_root, ant_label)
bee_train = AB(train_root, bee_label)
ant_test = AB(test_root, ant_label)
bee_test = AB(test_root, bee_label)
train_dataset = ant_train + bee_train
val_dataset = ant_test + bee_test
train_data = DataLoader(train_dataset, batch_size=16, shuffle=True)
val_data = DataLoader(val_dataset, batch_size=16, )
model = LeNet5().to(device)
# 设置损失函数、学习率、优化器
loss_fun = nn.CrossEntropyLoss().to(device)
learning_rate = 1e-3
optim = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9, weight_decay=1e-3)
# 接下来准备专题学习loss函数和优化器，如果有感兴趣的uu可以留下个关注以防走丢
epoch = 10
# 两层循环，实现训练
for i in range(epoch):
total_loss = 0 # 记录一个epoch的训练误差
for batch, data in enumerate(train_data):
img, label = data
img = img.to(device)
label = label.to(device)
model.train()
pred = model(img)
loss = loss_fun(pred, label)
total_loss = total_loss + loss
# 梯度清零、反向传播
optim.zero_grad()
loss.backward()
optim.step()
print("第{}个epoch的train_loss为：{}".format(i+1, total_loss))
writer.add_scalar("loss/traion", total_loss, i + 1)
torch.save(model.state_dict(), "./res/{}.pth".format(i + 1))
# 验证过程
val_loss = 0
model.eval()
with torch.no_grad():
correct = 0 # 分类正确率
for batch, data in enumerate(val_data):
img, label = data
img = img.to(device)
label = label.to(device)
pred = model(img)
loss = loss_fun(pred, label)
val_loss = val_loss + loss
correct += (pred.argmax(1) == label).type(torch.float).sum().item()
# print("第{}个epoch的val_loss为：{}".format(i + 1, val_loss))
print("第{}个epoch的精确度为：{}".format(i + 1, correct/len(val_dataset)))
# writer.add_scalar("loss/val", val_loss, i + 1)
print("done")总结
以上只是一次完整但极其简单的模型训练过程，有关数据的预处理、超参数选择等都没有涉及，以后都会分专题专门学习，如果您也刚开始学习深度学习，不妨留个关注，以后多多交流！
如果以上介绍有不完善或者错误的地方，还请uu们不吝赐教！