1. 导入模块
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 models import AlexNet2. 定义main函数
def main():3. 选择gpu/cpu
device = torch.device("cpu") # "cuda:0" if torch.cuda.is_available() else
print("using {} device.".format(device))4. 数据转换
# 数据转换
data_transform = {
"train": transforms.Compose([transforms.RandomResizedCrop(224), # RandomResizedCrop(224)随机裁剪为224*224
transforms.RandomHorizontalFlip(), # 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))])}RandomResizedCrop(224) 将图片随机裁剪为224*224
RandomHorizontalFlip() 以概率为默认值0.5将图片随机翻转
ToTensor() 将图像数据转化为张量模式，因为神经网络模型的输入是[N,C,H,W]
Normalize（） 是标准化处理，将图像数据的像素值缩放到[-1,1]的范围内。
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)使用pytorch中的datasets.ImageFolder类来创建一个图像文件夹数据集，同时应用之前的定义的数据转换对数据进行预处理。
6. 保存类别json文件
# {'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)class_to_idx是创建图像文件夹数据集自动创建的属性，用于将类别名称映射到类别索引
# {'daisy':0, 'dandelion':1, 'roses':2, 'sunflower':3, 'tulips':4}cla_dict是反转字典，将类别索引作为键，类别名称作值。
之后写入json格式文件。
{
"0": "daisy",
"1": "dandelion",
"2": "roses",
"3": "sunflowers",
"4": "tulips"
}7. 创建训练集和验证集的数据加载器（DataLoader)
batch_size = 8
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=True,
num_workers=nw)
print("using {} images for training, {} images for validation.".format(train_num,
val_num))nw是计算用于数据加载的工作进程数，将三个值中的最小值作为工作进程数，限制了并行加载数据的数量。
使用’torch.utils.data.DataLoader‘类创建数据加载器。其中train_dataset是之前创建的训练集数据集，batch_size参数指定每个批次的样本数量
8. 模型损失函数优化器等的设置。
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)创建一个AlexNet模型实例net，其中num_classes参数制定了模型的输出类别数，init_weights对模型的权重初始化。
CrossEntropyLoss（）是使用了交叉熵损失函数，适用于多分类问题。
optimizer 定义优化器，使用Adam优化器来更新模型参数，学习率设置为0.0002。
best_acc 最佳准确率，用于保持最佳模型
9. 训练神经网络模型的主要循环
for epoch in range(epochs):
# train net.train()方法的作用是将神经网络设置为训练模式。当调用net.train()时，PyTorch会将所有具有dropout和批量
# 归一化（Batch Normalization）等层设置为训练模式。 这里是为了防止过拟合，通过dropout方法随机失活一部分神经元
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() # 通过optimizer去更新每一个节点参数
# print statistics
running_loss += loss.item()
train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(epoch + 1,
epochs,
loss)
# validate
net.eval（) # 关闭dropout方法
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')net.train()将模型设置为训练模式，启用dropout和批量归一化等层的训练模式。dropout用于防止过拟合，训练过程中会随机失火一些神经元。
初始化runnning_loss为0.0。用于记录当前epoch的累计损失值。创建一个进度条train_bar,用于显示训练进度。
在每个epoch中，遍历训练数据加载器’train_loader'中的每个批次数据：
- 获取图像和标签数据
- 将梯度清零，以防止梯度累积
- 将图像数据传递给模型进行前向传播，得到模型的输出
- 计算损失函数
- 反向传播，计算梯度
- 更新模型参数
- 计算并累加当前批次的损失值
- 进度条中更新显示当前epoch的损失值
net.eval（) 在验证之前，将模型设置为评估模式，关闭dropout方法
使用torch.no_grad()上下文管理器，关闭梯度计算。
创建一个进度条val_bar,用于显示验证进度。
在每个验证集批次中，处理。最后计算验证机的准确率。
在最后，如果当前 epoch 的验证准确率超过之前记录的最佳准确率，则更新最佳准确率，并保存模型参数到指定路径。net.state_dict()方法，会返回一个字典，其中包含了模型的所有参数和对应的数值，这个字典的键是参数的名称，而值是包含了参数数值的张量，可以将这个状态字典保存到文件中，以便日后恢复模型的参数，或者将模型在其他地方重新创建。
