本文为B站UP 霹雳吧啦Wz 图片分类课程学习笔记，用于记录学习历程和个人复习
程序共分为三部分：model,train,predict。model.py用于存放模型，train.py用于存放训练时的程序，predict.py用于存放预测的程序，vgg16Net.pth为训练时val_accurate最高的模型参数，class_indices.json为类别索引映射。
注意：如要使用本文进行训练，需先按照 3.1 AlexNet网络结构详解与花分类数据集下载 在指定位置放好图片数据集方可进行训练。
个人感觉比较好的学习办法是先听一遍up的讲解，再看代码，哪段代码不会就查哪里，把所有的代码都搞清楚再自己改改代码。
model.py
总体代码
import torch.nn as nn
import torch
# official pretrain weights
model_urls = {
'vgg11': 'https://download.pytorch.org/models/vgg11-bbd30ac9.pth',
'vgg13': 'https://download.pytorch.org/models/vgg13-c768596a.pth',
'vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth',
'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth'
}
class VGG(nn.Module):
def __init__(self, features, num_classes=1000, init_weights=False):
super(VGG, self).__init__()
self.features = features
self.classifier = nn.Sequential(
nn.Linear(512*7*7, 4096),
nn.ReLU(True),
nn.Dropout(p=0.5),
nn.Linear(4096, 4096),
nn.ReLU(True),
nn.Dropout(p=0.5),
nn.Linear(4096, num_classes)
)
if init_weights:
self._initialize_weights()
def forward(self, x):
# N x 3 x 224 x 224
x = self.features(x)
# N x 512 x 7 x 7
x = torch.flatten(x, start_dim=1)
# N x 512*7*7
x = self.classifier(x)
return x
def _initialize_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
# nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
nn.init.xavier_uniform_(m.weight)
if m.bias is not None: #查看是否有偏置项
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.Linear):
nn.init.xavier_uniform_(m.weight)
# nn.init.normal_(m.weight, 0, 0.01)
nn.init.constant_(m.bias, 0)
def make_features(cfg: list):
layers = []
in_channels = 3
for v in cfg:
if v == "M":
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
else:
conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
layers += [conv2d, nn.ReLU(True)]
in_channels = v
return nn.Sequential(*layers)
cfgs = {
'vgg11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
'vgg13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
'vgg16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
'vgg19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
}
def vgg(model_name="vgg16", **kwargs):
assert model_name in cfgs, "Warning: model number {} not in cfgs dict!".format(model_name)
cfg = cfgs[model_name]
model = VGG(make_features(cfg), **kwargs)
return model
train.py
数据预处理
import os
import sys
import json
import torch
import torch.nn as nn
from torchvision import transforms, datasets
import torch.optim as optim
from tqdm import tqdm
from model import vgg
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("using {} device.".format(device))
data_transform = {
"train": transforms.Compose([transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]),
"val": transforms.Compose([transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])}
这段代码定义了图像数据的预处理步骤，这些步骤在训练和验证过程中对图像进行处理。
data_transform 字典
这个字典包含两个键："train" 和 "val"，分别对应训练和验证数据的预处理步骤。
训练数据的预处理 ("train")
transforms.Compose([
transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ])
transforms.RandomResizedCrop(224): 这个函数随机裁剪图像，生成一个大小为 224x224 的图像区域。这有助于模型学习到图像的不同部分，增加模型的泛化能力。transforms.RandomHorizontalFlip(): 这个函数随机地水平翻转图像。这同样有助于增加模型的泛化能力，因为图像的方向在实际应用中可能会变化。transforms.ToTensor(): 将 PIL 图像或 NumPyndarray转换为FloatTensor，并将图像的像素值从 [0, 255] 归一化到 [0.0, 1.0]。transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)): 将图像的每个通道的像素值进行标准化，使其均值为 0.5，标准差为 0.5。这有助于模型训练时的数值稳定性。
验证数据的预处理 ("val")
transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ])
transforms.Resize((224, 224)): 将图像调整到固定的大小 224x224。这确保了所有验证图像都具有相同的尺寸，这对于模型的一致性输入是必要的。transforms.ToTensor(): 与训练数据相同，将图像转换为FloatTensor，并归一化像素值。transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)): 与训练数据相同，进行像素值的标准化。- 为什么使用这些预处理步骤？
- 随机性：通过随机裁剪和随机翻转，模型能够学习到图像的不同变体，增强其对新数据的适应能力。
- 标准化：将图像像素值归一化到 [0, 1] 或 [0.5, 0.5] 的范围内，有助于模型更快地收敛，并减少训练过程中的数值不稳定性。
- 一致性：确保训练和验证图像具有相同的尺寸和数值范围，有助于模型在不同阶段的一致性学习。
这些预处理步骤是图像处理和深度学习中常见的实践，有助于提高模型的性能和泛化能力。
图片路径准备
data_root = os.path.abspath(os.path.join(os.getcwd(), "../..")) # get data root path
image_path = os.path.join(data_root, "data_set", "flower_data") # flower data set path
assert os.path.exists(image_path), "{} path does not exist.".format(image_path)
这段代码用于设置数据集的根路径，并确保该路径存在。
-
获取当前工作目录：
data_root = os.path.abspath(os.path.join(os.getcwd(), "../.."))os.getcwd()获取当前工作目录的路径。os.path.join(os.getcwd(), "../..")将当前工作目录向上移动两个目录层级。这通常用于从当前脚本所在的目录跳转到项目的根目录。os.path.abspath()将相对路径转换为绝对路径，确保路径的准确性和可移植性。
-
设置数据集路径：
image_path = os.path.join(data_root, "data_set", "flower_data")os.path.join(data_root, "data_set", "flower_data")将根目录与子目录 "data_set" 和 "flower_data" 连接起来，形成完整的数据集路径。这里假设数据集存储在项目的 "data_set" 文件夹下的 "flower_data" 文件夹中。
-
检查路径是否存在：
assert os.path.exists(image_path),"{} path does not exist.".format(image_path)os.path.exists(image_path)检查image_path指定的路径是否存在。assert是一个断言语句，用于确保某个条件为真。如果条件为假，则抛出异常。- 如果
image_path指定的路径不存在，将抛出一个AssertionError异常，并显示错误消息"{path} path does not exist."，其中{path}会被替换为实际的路径。
代码的作用：
- 确定根目录：通过向上移动两个目录层级，确保脚本可以从项目的任何子目录运行，而不仅仅是项目根目录。
- 构建数据集路径：通过连接根目录和数据集子目录，构建出完整的数据集路径。
- 确保路径有效：通过断言检查，确保数据集路径存在，避免在后续操作中因路径错误而导致程序崩溃。
这种路径设置方式常用于项目结构较为复杂的场景，其中数据集、模型和其他资源可能分散在不同的目录中。通过这种方式，可以确保脚本在任何位置运行时都能正确地找到所需的资源。
加载图片
train_dataset = datasets.ImageFolder(root=os.path.join(image_path, "train"),
transform=data_transform["train"])
train_num = len(train_dataset)
validate_dataset = datasets.ImageFolder(root=os.path.join(image_path, "val"),
transform=data_transform["val"])
val_num = len(validate_dataset)datasets.ImageFolder：这是 PyTorch 提供的一个类，用于从文件夹中加载图像数据。每个子文件夹代表一个类别，文件夹名称是类别的名称。root=os.path.join(image_path, "train")：指定数据集的根目录。这里使用os.path.join将之前定义的image_path和子目录"train"连接起来，形成完整的训练数据集路径。transform=data_transform["train"]：应用之前定义的训练数据预处理步骤。这些步骤包括随机裁剪、随机水平翻转、转换为张量以及归一化。
保存类别索引映射
flower_list = train_dataset.class_to_idx
cla_dict = dict((val, key) for key, val in flower_list.items())
json_str = json.dumps(cla_dict, indent=4)
with open('class_indices.json', 'w') as json_file:
json_file.write(json_str)这段代码的作用是将训练数据集中的类别索引映射保存到一个 JSON 文件中。
-
获取类别索引映射：
flower_list = train_dataset.class_to_idxtrain_dataset.class_to_idx是一个字典，其中键是类别的名称（如'daisy'），值是该类别在数据集中的索引（如0）。这个映射是由ImageFolder在加载数据集时自动生成的。 -
反转类别索引映射：
这行代码将cla_dict = dict((val, key) for key, val in flower_list.items())flower_list字典的键和值对调，生成一个新的字典cla_dict。在这个新字典中，索引（如0）是键，类别名称（如'daisy'）是值。这样做的目的是方便在模型预测时将预测的索引转换回类别名称。 -
将字典转换为 JSON 字符串：
json_str = json.dumps(cla_dict, indent=4)json.dumps()函数将 Python 字典转换为 JSON 格式的字符串。indent=4参数使生成的 JSON 字符串具有可读性，即每级缩进 4 个空格。 -
将 JSON 字符串写入文件：
使用with open('class_indices.json', 'w') as json_file: json_file.write(json_str)with语句打开文件'class_indices.json'进行写操作。'w'模式表示如果文件已存在，则会被覆盖；如果文件不存在，则会被创建。json_file.write(json_str)将 JSON 字符串写入文件。
代码的作用：
- 保存类别索引映射：通过将类别名称和索引的映射保存到文件，可以在模型训练和预测时方便地进行类别名称和索引之间的转换。
- 便于模型预测：在模型预测时，通常需要将预测的索引转换回类别名称，以便更容易地理解预测结果。
设置批量大小
batch_size = 32
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8]) # number of workers
print('Using {} dataloader workers every process'.format(nw))[os.cpu_count(), batch_size if batch_size > 1 else 0, 8]
这个是一个列表，由三部分组成：
os.cpu_count()返回当前系统可用的 CPU 核心数。batch_size if batch_size > 1 else 0检查批量大小是否大于 1，如果是，则使用批量大小作为工作进程数；如果不是（例如批量大小为 1），则不使用额外的工作进程。min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])从 CPU 核心数、调整后的工作进程数（基于批量大小）和固定值 8 中选择最小的一个，作为最终的工作进程数。这样做是为了避免使用过多的资源，尤其是在 CPU 核心数较少或批量大小较小的情况下。- 限制工作进程数的最大值为 8 是一个常见的做法，以确保不会因过多的并行进程而过度消耗系统资源。
创建数据加载器
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size, shuffle=True,
num_workers=nw)
validate_loader = torch.utils.data.DataLoader(validate_dataset,
batch_size=batch_size, shuffle=False,
num_workers=nw)
print("using {} images for training, {} images for validation.".format(train_num,
val_num))
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size, shuffle=True, num_workers=nw)
torch.utils.data.DataLoader：这是 PyTorch 提供的一个类，用于封装数据集并提供批量加载和多进程加载功能。train_dataset：这是之前创建的训练数据集对象。batch_size=batch_size：设置每个批次的样本数量。这里使用之前定义的batch_size变量，其值为 32。shuffle=True：在每个 epoch 开始时，是否对数据进行随机打乱。这有助于模型训练时的泛化能力，避免模型对数据的特定顺序产生依赖。num_workers=nw：设置用于数据加载的工作进程数。这里使用之前计算的nw变量，其值是os.cpu_count()、batch_size和 8 中的最小值。这有助于提高数据加载的效率。
准备训练参数
model_name = "vgg16"
net = vgg(model_name=model_name, num_classes=5, init_weights=True)
net.to(device)
loss_function = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.0001)
epochs = 6
best_acc = 0.5
save_path = './{}Net.pth'.format(model_name)
train_steps = len(train_loader)
save_path = './{}Net.pth'.format(model_name)
这行代码是用于定义模型权重文件的保存路径。
-
./：这是一个相对路径，表示当前工作目录。这意味着文件将被保存在运行脚本的当前目录中。 -
{}Net.pth：这是一个格式化字符串，花括号{}是一个占位符，用于插入变量值。 -
format(model_name)：这是str.format()方法，用于将model_name变量的值替换到格式化字符串中的占位符{}位置。model_name变量在之前的代码中被设置为"vgg16"。 -
Net.pth：这是一个文件名后缀，通常用于表示 PyTorch 模型权重文件。.pth是一个常用的文件扩展名，用于区分 PyTorch 模型文件。
开始训练
for epoch in range(epochs):
# train
net.train() # 设置为训练模式
running_loss = 0.0
train_bar = tqdm(train_loader, file=sys.stdout)
for step, data in enumerate(train_bar):
images, labels = data
optimizer.zero_grad()
outputs = net(images.to(device))
loss = loss_function(outputs, labels.to(device))
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(epoch + 1, epochs, loss)
# validate
net.eval()
acc = 0.0 # accumulate accurate number / epoch
with torch.no_grad():
val_bar = tqdm(validate_loader, file=sys.stdout)
for val_data in val_bar:
val_images, val_labels = val_data
outputs = net(val_images.to(device))
predict_y = torch.max(outputs, dim=1)[1]
acc += torch.eq(predict_y, val_labels.to(device)).sum().item()
val_accurate = acc / val_num
print('[epoch %d] train_loss: %.3f val_accuracy: %.3f' %
(epoch + 1, running_loss / train_steps, val_accurate))
if val_accurate > best_acc:
best_acc = val_accurate
torch.save(net.state_dict(), save_path)
print('Finished Training')使用 tqdm 显示进度
train_bar = tqdm(train_loader, file=sys.stdout)
tqdm是一个进度条库，用于显示训练过程中的进度。train_loader是数据加载器，负责按批次提供训练数据。file=sys.stdout指定进度条的输出位置，这里是标准输出（即控制台）。
内部循环：处理每个批次的数据
for step, data in enumerate(train_bar):
enumerate(train_bar)遍历train_bar，step是批次的索引，data是当前批次的数据。- 每个
data包含一对图像和标签，通常是通过train_loader提供的。
数据处理
images, labels = data
optimizer.zero_grad()
outputs = net(images.to(device))
loss = loss_function(outputs, labels.to(device))
loss.backward()
optimizer.step()
images, labels = data：从data中解包出图像和标签。optimizer.zero_grad()：清除之前的梯度，为新的反向传播准备。这是每次迭代开始时必须执行的步骤。outputs = net(images.to(device))：将图像数据移动到指定的设备（如 GPU），并通过模型前向传播得到输出。loss = loss_function(outputs, labels.to(device))：计算输出和标签之间的损失。这里使用的是交叉熵损失函数loss_function。loss.backward()：反向传播损失，计算模型参数的梯度。optimizer.step()：根据计算出的梯度更新模型的权重。
更新运行损失并显示进度
running_loss += loss.item()
train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(epoch + 1, epochs, loss)
running_loss += loss.item()：将当前批次的损失添加到运行损失中。train_bar.desc：更新进度条的描述，显示当前 epoch、总 epoch 数和当前批次的损失。
