train.py的全部代码如下：
python
import os
import sys
import json
import torch
import torch.nn as nn
from torchvision import transforms, datasets, utils
import matplotlib.pyplot as plt
import numpy as np
import torch.optim as optim
from tqdm import tqdm
from model import AlexNet
def main():
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)), # cannot 224, must (224, 224)
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 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)
train_dataset = datasets.ImageFolder(root=os.path.join(image_path, "train"),
transform=data_transform["train"])
train_num = len(train_dataset)
# {'daisy':0, 'dandelion':1, 'roses':2, 'sunflower':3, 'tulips':4}
flower_list = train_dataset.class_to_idx
cla_dict = dict((val, key) for key, val in flower_list.items())
# write dict into json file
json_str = json.dumps(cla_dict, indent=4)
with open('class_indices.json', 'w') as json_file:
json_file.write(json_str)
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))
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size, shuffle=True,
num_workers=nw)
validate_dataset = datasets.ImageFolder(root=os.path.join(image_path, "val"),
transform=data_transform["val"])
val_num = len(validate_dataset)
validate_loader = torch.utils.data.DataLoader(validate_dataset,
batch_size=4, shuffle=False,
num_workers=nw)
print("using {} images for training, {} images for validation.".format(train_num,
val_num))
# test_data_iter = iter(validate_loader)
# test_image, test_label = test_data_iter.next()
#
# def imshow(img):
# img = img / 2 + 0.5 # unnormalize
# npimg = img.numpy()
# plt.imshow(np.transpose(npimg, (1, 2, 0)))
# plt.show()
#
# print(' '.join('%5s' % cla_dict[test_label[j].item()] for j in range(4)))
# imshow(utils.make_grid(test_image))
net = AlexNet(num_classes=5, init_weights=True)
net.to(device)
loss_function = nn.CrossEntropyLoss()
# pata = list(net.parameters())
optimizer = optim.Adam(net.parameters(), lr=0.0002)
epochs = 10
save_path = './AlexNet.pth'
best_acc = 0.0
train_steps = len(train_loader)
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')
if __name__ == '__main__':
main()import os: 导入Python的操作系统模块，用于与操作系统交互。
import sys: 导入Python的系统模块，提供对Python解释器使用或维护的一些变量的访问，以及与Python解释器强烈交互的功能。
import json: 导入Python的JSON模块，用于处理JSON格式的数据。
import torch: 导入PyTorch库，一个开源的深度学习框架。
import torch.nn as nn: 从PyTorch库中导入神经网络模块，并简称为nn。
from torchvision import transforms, datasets, utils: 从torchvision库中导入transforms（数据预处理）、datasets（数据集）和utils（实用工具）模块。
import matplotlib.pyplot as plt: 导入matplotlib的pyplot模块，用于绘制图形。
import numpy as np: 导入NumPy库，并简称为np，用于进行数值计算。
import torch.optim as optim: 从PyTorch库中导入优化器模块，并简称为optim。
from tqdm import tqdm: 从tqdm库中导入tqdm模块，用于显示进度条。
from model import AlexNet: 从当前目录下的model模块中导入AlexNet类
python
def main():
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)), # cannot 224, must (224, 224)
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 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)
train_dataset = datasets.ImageFolder(root=os.path.join(image_path, "train"),
transform=data_transform["train"])
train_num = len(train_dataset)device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
创建一个设备对象，如果CUDA（一种用于运行深度学习模型的并行计算平台和API模型）可用，则使用GPU（这里是"cuda:0"表示第一个GPU），否则使用CPU。
print("using {} device.".format(device))
打印出正在使用的设备（CPU或GPU）。
data_transform = {
定义一个字典data_transform，用于存储训练和验证数据集的预处理步骤。
"train": transforms.Compose([transforms.RandomResizedCrop(224),
对于训练数据集，首先进行随机大小和比例的裁剪（RandomResizedCrop），裁剪的大小为224x224像素。
transforms.RandomHorizontalFlip(),
接着，进行随机水平翻转，这是一个数据增强的操作，目的是增加模型的泛化能力。
transforms.ToTensor(),
将PIL图像或NumPy ndarray转换为PyTorch张量。
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]),
对图像进行标准化，使其具有零均值和单位方差。
"val": transforms.Compose([transforms.Resize((224, 224)),
对于验证数据集，首先调整图像大小为224x224像素。注意这里明确指出了宽度和高度，与"train"的处理略有不同。
transforms.ToTensor(),
同上，将图像转换为PyTorch张量。
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])}
同上，对图像进行标准化。
data_root = os.path.abspath(os.path.join(os.getcwd(), ".../..."))
获取数据根路径，通常是相对于当前工作目录的上一级目录。这通常用于组织数据集和其他资源。
image_path = os.path.join(data_root, "data_set", "flower_data")
定义图片数据集的路径，它位于之前获取的数据根目录下的"data_set"文件夹中，并命名为"flower_data"。
assert os.path.exists(image_path), "{} path does not exist.".format(image_path)
使用断言确保图片路径存在。如果路径不存在，程序将抛出一个错误。
train_dataset = datasets.ImageFolder(root=os.path.join(image_path, "train"), transform=data_transform["train"])
使用ImageFolder从指定路径加载训练数据集，并应用之前定义的训练预处理。这假设文件夹结构中每个子文件夹名表示一个类别。
train_num = len(train_dataset)
获取训练数据集的大小（即图像数量）。
python
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))
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size, shuffle=True,
num_workers=nw)
validate_dataset = datasets.ImageFolder(root=os.path.join(image_path, "val"),
transform=data_transform["val"])
val_num = len(validate_dataset)
validate_loader = torch.utils.data.DataLoader(validate_dataset,
batch_size=4, shuffle=False,
num_workers=nw)
print("using {} images for training, {} images for validation.".format(train_num,
val_num))batch_size = 32
定义批量大小为32，这是在训练和验证过程中一次处理的数据样本数量。
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])
计算数据加载的工作进程数（workers）。
os.cpu_count() 返回计算机的CPU核心数。
batch_size if batch_size > 1 else 0 表示如果批量大小大于1，则使用批量大小作为工作进程数，否则为0。
最后，工作进程数被限制为8。这是为了防止数据加载时使用过多的系统资源。
print('Using {} dataloader workers every process'.format(nw))
打印出使用的工作进程数。
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size, shuffle=True,
num_workers=nw)
使用DataLoader从train_dataset中加载训练数据。
batch_size=batch_size 设置批量大小。
shuffle=True 表示在每个训练时代开始时打乱数据顺序。
num_workers=nw 设置数据加载的工作进程数。
transform=data_transform["val"])
val_num = len(validate_dataset)
validate_loader = torch.utils.data.DataLoader(validate_dataset,
batch_size=4, shuffle=False,
num_workers=nw)
使用DataLoader从validate_dataset中加载验证数据。
batch_size=4 设置批量大小为4。验证时通常使用较小的批量大小。
shuffle=False 表示验证时不打乱数据顺序。
num_workers=nw 设置数据加载的工作进程数。
python
net = AlexNet(num_classes=5, init_weights=True)
net.to(device)
loss_function = nn.CrossEntropyLoss()
# pata = list(net.parameters())
optimizer = optim.Adam(net.parameters(), lr=0.0002)
epochs = 10
save_path = './AlexNet.pth'
best_acc = 0.0
train_steps = len(train_loader)
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)net = AlexNet(num_classes=5, init_weights=True) 实例化一个AlexNet网络
这里设置分类是5
net.to(device) 把模型移动到指定的设备上
loss_function = nn.CrossEntropyLoss() 指定损失函数为交叉熵损失函数
optimizer = optim.Adam(net.parameters(), lr=0.0002) 定义了Adam优化器和学习率
epochs = 10 # 训练轮数
save_path = './AlexNet.pth' # 模型保存路径
best_acc = 0.0 # 当前最高准确率
train_steps = len(train_loader) # 训练批次总数
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)
模型设置为训练模型
初始化运行损失为0 使用进度条显示当前的训练状态
对于每个训练epoch
获取批次数据和标签。
将数据和标签移动到指定设备。
清除优化器的梯度。
前向传播。
计算损失。
反向传播。
优化器进行参数更新。
更新运行损失。
更新进度条描述。
python
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)1.设置模型为评估模式
net.eval()将模型设置为评估模式，这会关闭如dropout和batch normalization这类只在训练时使用的层
- 初始化准确率变量:
acc = 0.0 # accumulate accurate number / epoch
初始化一个变量来累计每个epoch的正确预测数。
3.开始验证循环: 使用torch.no_grad()上下文管理器来关闭自动梯度计算，因为我们只关心验证结果，不需要计算梯度。
对于每个验证批次：
获取批次数据和标签。
将数据和标签移动到指定设备。
通过模型进行前向传播。
使用torch.max获取预测的类别。
检查预测类别是否与真实标签相同。
累计正确预测的数量。
4.计算验证集上的准确率。
-
打印训练和验证结果: 使用字符串格式化打印当前的训练损失和验证准确率。
-
更新最佳准确率并保存模型: 如果当前准确率高于之前的最高准确率，则更新最高准确率并将模型保存到指定路径。