【无标题】
代码中需要注意的是：
pytorch 中 Tensor 参数的顺序为 (batch, channel, height, width) ，下面代码中没有写batch
卷积的参数为Conv2d(in_channels, out_channels, kernel_size, stride, padding, ...)，一般关心这5个参数即可
卷积池化层提取图像特征，全连接层进行图像分类，代码中写成两个模块，方便调用(nn.Sequential)
为了加快训练，代码只使用了一半的网络参数，相当于只用了原论文中网络结构的下半部分（正好原论文中用的双GPU，我的电脑只有一块GPU）（后来我又用完整网络跑了遍，发现一半参数跟完整参数的训练结果acc相差无几）
import torch.nn as nn
import torch
class AlexNet(nn.Module):
def __init__(self, num_classes=1000, init_weights=False):
super(AlexNet, self).__init__()
# 用nn.Sequential()将网络打包成一个模块，精简代码
self.features = nn.Sequential( # 卷积层提取图像特征
nn.Conv2d(3, 48, kernel_size=11, stride=4, padding=2), # input[3, 224, 224] output[48, 55, 55]
nn.ReLU(inplace=True), # 直接修改覆盖原值，节省运算内存
nn.MaxPool2d(kernel_size=3, stride=2), # output[48, 27, 27]
nn.Conv2d(48, 128, kernel_size=5, padding=2), # output[128, 27, 27]
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2), # output[128, 13, 13]
nn.Conv2d(128, 192, kernel_size=3, padding=1), # output[192, 13, 13]
nn.ReLU(inplace=True),
nn.Conv2d(192, 192, kernel_size=3, padding=1), # output[192, 13, 13]
nn.ReLU(inplace=True),
nn.Conv2d(192, 128, kernel_size=3, padding=1), # output[128, 13, 13]
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2), # output[128, 6, 6]
)
self.classifier = nn.Sequential( # 全连接层对图像分类
nn.Dropout(p=0.5), # Dropout 随机失活神经元，默认比例为0.5
nn.Linear(128 * 6 * 6, 2048),
nn.ReLU(inplace=True),
nn.Dropout(p=0.5),
nn.Linear(2048, 2048),
nn.ReLU(inplace=True),
nn.Linear(2048, num_classes), # 输出是数据集的类别个数 num_classes
)
if init_weights: # 初始化权重的选项
self._initialize_weights()
# 前向传播过程
def forward(self, x): # x为输入样本
x = self.features(x) # x输入features的一系列结构
x = torch.flatten(x, start_dim=1) # 展平后再传入全连接层 start_dim=1代表从channel维度开始
x = self.classifier(x)
return x
# 网络权重初始化，实际上 pytorch 在构建网络时会自动初始化权重
def _initialize_weights(self):
for m in self.modules(): # 遍历self.modules(),继承自nn.Module的
if isinstance(m, nn.Conv2d): # 若m结构是卷积层
nn.init.kaiming_normal_(m.weight, mode='fan_out', # 用（何）kaiming_normal_法初始化权重weight
nonlinearity='relu')
if m.bias is not None: # 若偏置不为空
nn.init.constant_(m.bias, 0) # 初始化偏重为0
elif isinstance(m, nn.Linear): # 若是全连接层
nn.init.normal_(m.weight, 0, 0.01) # 正态分布初始化
nn.init.constant_(m.bias, 0) # 初始化偏重为0
nn.Conv2d(3, 48, kernel_size=11, stride=4, padding=2)
padding只能传入两种类型。一种是int整型，如padding=1，则在特征矩阵的上下方补一层0，左右补一列0；一种是tuple，如tuple(1,2),则1代表上下各补一行0, 2代表左右两侧各补两列0。
nn.ZeroPad2d((1,2,1,2))
左侧补一列，右侧补两列
上方补一行，下方补两行
2. train.py
# 导入包
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 model import AlexNet
import os
import json
import time
# 使用GPU训练
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)
2.1 数据预处理
对训练集的预处理，多了随机裁剪和水平翻转这两个步骤。可以起到扩充数据集的作用，增强模型泛化能力。
(40条消息) pytorch中的transforms.ToTensor和transforms.Normalize理解_秃头小苏的博客-CSDN博客
Pytorch中transforms.RandomResizedCrop()等图像操作_心向林丶的博客-CSDN博客_transforms.randomcrop
data_transform = {
"train": transforms.Compose([transforms.RandomResizedCrop(224), # 随机裁剪，再缩放成 224×224
transforms.RandomHorizontalFlip(p=0.5), # 水平方向随机翻转，概率为 0.5, 即一半的概率翻转, 一半的概率不翻转
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))])}
2.2 导入、加载 训练集
还记得上篇用到的 CIFAR10 是怎么导入和加载数据集的吗？
上篇LeNet网络搭建中是使用的torchvision.datasets.CIFAR10和torch.utils.data.DataLoader()来导入和加载数据集。
# 导入训练集
train_set = torchvision.datasets.CIFAR10(root='./data', # 数据集存放目录
train=True, # 表示是数据集中的训练集
download=True, # 第一次运行时为True，下载数据集，下载完成后改为False
transform=transform) # 预处理过程
# 加载训练集
train_loader = torch.utils.data.DataLoader(train_set, # 导入的训练集
batch_size=50, # 每批训练的样本数
shuffle=False, # 是否打乱训练集
num_workers=0) # num_workers在windows下设置为0
但是这次的 花分类数据集 并不在 pytorch 的 torchvision.datasets. 中，因此需要用到datasets.ImageFolder()来导入。
ImageFolder()返回的对象是一个包含数据集所有图像及对应标签构成的二维元组容器，支持索引和迭代，可作为torch.utils.data.DataLoader的输入。具体可参考：pytorch ImageFolder和Dataloader加载自制图像数据集
# 获取图像数据集的路径
data_root = os.path.abspath(os.path.join(os.getcwd(), "../..")) # get data root path 返回上上层目录
image_path = data_root + "/data_set/flower_data/" # flower data_set path
# 导入训练集并进行预处理
train_dataset = datasets.ImageFolder(root=image_path + "/train",
transform=data_transform["train"])
train_num = len(train_dataset) # 训练集有多少张图片
# 按batch_size分批次加载训练集
train_loader = torch.utils.data.DataLoader(train_dataset, # 导入的训练集
batch_size=32, # 每批训练的样本数
shuffle=True, # 是否打乱训练集
num_workers=0) # 使用线程数，在windows下设置为0
os.getcwd() 获取当前文件所在的目录 ''..''代表返回上一层目录 ''../..''代表返回上上层目录
os.path.join 将后面两个路径连接在一起
2.3 导入、加载 验证集
# 导入验证集并进行预处理
validate_dataset = datasets.ImageFolder(root=image_path + "/val",
transform=data_transform["val"])
val_num = len(validate_dataset)
# 加载验证集
validate_loader = torch.utils.data.DataLoader(validate_dataset, # 导入的验证集
batch_size=32,
shuffle=True,
num_workers=0)
2.4 存储 索引：标签 的字典
# 字典，类别：索引 {'daisy':0, 'dandelion':1, 'roses':2, 'sunflower':3, 'tulips':4}
flower_list = train_dataset.class_to_idx
# 将 flower_list 中的 key 和 val 调换位置
cla_dict = dict((val, key) for key, val in flower_list.items())
# 将 cla_dict 写入 json 文件中
json_str = json.dumps(cla_dict, indent=4)
with open('class_indices.json', 'w') as json_file:
json_file.write(json_str)
class_indices.json 文件内容如下：
{
"0": "daisy",
"1": "dandelion",
"2": "roses",
"3": "sunflowers",
"4": "tulips"
}
2.5 训练过程
训练过程中需要注意：
net.train()：训练过程中开启 Dropoutnet.eval()： 验证过程关闭 Dropout
net = AlexNet(num_classes=5, init_weights=True) # 实例化网络（输出类型为5，初始化权重）
net.to(device) # 分配网络到指定的设备（GPU/CPU）训练
loss_function = nn.CrossEntropyLoss() # 交叉熵损失
optimizer = optim.Adam(net.parameters(), lr=0.0002) # 优化器（训练参数，学习率）
save_path = './AlexNet.pth'
best_acc = 0.0 # 历史最优准确率
for epoch in range(10):
########################################## train ###############################################
net.train() # 训练过程中开启 Dropout
running_loss = 0.0 # 每个 epoch 都会对 running_loss 清零
time_start = time.perf_counter() # 对训练一个 epoch 计时
for step, data in enumerate(train_loader, start=0): # 遍历训练集，step从0开始计算
images, labels = data # 获取训练集的图像和标签
optimizer.zero_grad() # 清除历史梯度
outputs = net(images.to(device)) # 正向传播
loss = loss_function(outputs, labels.to(device)) # 计算损失
loss.backward() # 反向传播 到节点中
optimizer.step() # 优化器更新参数
running_loss += loss.item()
# 打印训练进度（使训练过程可视化）
rate = (step + 1) / len(train_loader) # 当前进度 = 当前step / 训练一轮epoch所需总step
a = "*" * int(rate * 50)
b = "." * int((1 - rate) * 50)
print("\rtrain loss: {:^3.0f}%[{}->{}]{:.3f}".format(int(rate * 100), a, b, loss), end="")
print()
print('%f s' % (time.perf_counter()-time_start))
########################################### validate ###########################################
net.eval() # 验证过程中关闭 Dropout
acc = 0.0 # accumulate accurate number / epoch
with torch.no_grad():
for val_data in validate_loader:
val_images, val_labels = val_data
outputs = net(val_images.to(device)) # 验证图片指认到device上，传到网络，进行正向传播，得到输出
predict_y = torch.max(outputs, dim=1)[1] # 以output中值最大位置对应的索引（标签）作为预测输出
acc += (predict_y == val_labels.to(device)).sum().item()
val_accurate = acc / val_num # 验证集正确率 = 验证正确样本个数/总的样本数
# 保存准确率最高的那次网络参数
if val_accurate > best_acc: # 当前准确率大于历史最优的准确率
best_acc = val_accurate
torch.save(net.state_dict(), save_path)
print('[epoch %d] train_loss: %.3f test_accuracy: %.3f \n' %
(epoch + 1, running_loss / step, val_accurate))
print('Finished Training')
训练打印信息如下：
cuda
train loss: 100%[**************************************************->]1.566
27.450399 s
[epoch 1] train_loss: 1.413 test_accuracy: 0.404
train loss: 100%[**************************************************->]1.412
27.897467399999996 s
[epoch 2] train_loss: 1.211 test_accuracy: 0.503
train loss: 100%[**************************************************->]1.412
28.665594 s
[epoch 3] train_loss: 1.138 test_accuracy: 0.544
train loss: 100%[**************************************************->]0.924
28.6858524 s
[epoch 4] train_loss: 1.075 test_accuracy: 0.621
train loss: 100%[**************************************************->]1.200
28.020624199999986 s
[epoch 5] train_loss: 1.009 test_accuracy: 0.621
train loss: 100%[**************************************************->]0.985
27.973145999999986 s
[epoch 6] train_loss: 0.948 test_accuracy: 0.607
train loss: 100%[**************************************************->]0.583
28.290610200000003 s
[epoch 7] train_loss: 0.914 test_accuracy: 0.670
train loss: 100%[**************************************************->]0.930
28.51416950000001 s
[epoch 8] train_loss: 0.912 test_accuracy: 0.621
train loss: 100%[**************************************************->]1.210
28.98158360000002 s
[epoch 9] train_loss: 0.840 test_accuracy: 0.668
train loss: 100%[**************************************************->]0.961
28.330670499999997 sp
[epoch 10] train_loss: 0.833 test_accuracy: 0.684
Finished Training
完整代码
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()
3. predict.py
import torch
from model import AlexNet
from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as plt
import json
# 预处理
data_transform = transforms.Compose(
[transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# load image
img = Image.open("蒲公英.jpg")
plt.imshow(img)
# [N, C, H, W]
img = data_transform(img)
# expand batch dimension
img = torch.unsqueeze(img, dim=0)
# read class_indict
try:
json_file = open('./class_indices.json', 'r')
class_indict = json.load(json_file)
except Exception as e:
print(e)
exit(-1)
# create model
model = AlexNet(num_classes=5)
# load model weights
model_weight_path = "./AlexNet.pth"
model.load_state_dict(torch.load(model_weight_path))
# 关闭 Dropout
model.eval()
with torch.no_grad():
# predict class
output = torch.squeeze(model(img)) # 将输出压缩，即压缩掉 batch 这个维度
predict = torch.softmax(output, dim=0)
predict_cla = torch.argmax(predict).numpy()
print(class_indict[str(predict_cla)], predict[predict_cla].item())
plt.show()
打印出预测的标签以及概率值：
dandelion 0.7221569418907166