使用pytorch搭建AlexNet并训练花分类数据集
深度学习学习笔记
导师博客：https://blog.csdn.net/qq_37541097/article/details/103482003
导师github:https://github.com/WZMIAOMIAO/deep-learning-for-image-processing
代码用的导师的，自己又加了些备注，就放在自己的github里了：
https://github.com/Petrichor223/Deep_Learning/tree/master
网络结构
网络介绍及结构这一部分导师写的很详细：AlexNet网络结构详解与模型的搭建
数据集介绍
花分类数据集下载：http://download.tensorflow.org/example_images/flower_photos.tgz
训练集和测试集的划分：
参考：https://github.com/WZMIAOMIAO/deep-learning-for-image-processing/tree/master/data_set
下载完成后将压缩包放在data_set文件夹下并解压到新建的flower_data文件夹下，然后再data_set文件夹下按住shift + 右键打开Powershell窗口，在窗口中输入python .\split_data.py去运行脚本，就会按照9：1进行数据集的划分
划分完成后在flower_data文件夹下就会有生成的训练集和验证集
项目文件
AlexNet
├─ model.py
├─ predict.py
├─ train.py
├─ flower_data
1. model.py
Alexnet所有层的参数表：
import torch.nn as nn
import torch
class AlexNet(nn.Module):
def __init__(self, num_classes=1000, init_weights=False): #初始化函数来定义网络在正向传播过程中使用的层结构
super(AlexNet, self).__init__()
# n.Sequential能够将一系列的层结构进行打包组合成新的结构，如果像LeNet那样定义每一层会非常麻烦
self.features = nn.Sequential( #这里对应特征提取
#对照Alexnet所有层参数表进行设配置
nn.Conv2d(3, 48, kernel_size=11, stride=4, padding=2), # input[3, 224, 224] output[48, 55, 55] 由于花分类数据集较小，将卷积核个数减半，为了方便直接将padding设置为2
nn.ReLU(inplace=True), #inplace可以理解为pytorch增加计算量但降低内存使用
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), #一般加载全连接层之间，默认失活比例为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),
)
if init_weights: #初始化权重
self._initialize_weights()
def forward(self, x): #正向传播过程
x = self.features(x) #首先将输入的训练样本x输入到features中
x = torch.flatten(x, start_dim=1) #将上一步的输出进行展平处理，维度从dim=1开始(0维度为batch，1维度为channel)
x = self.classifier(x) #展平后将其输入到分类结构中
return x
# 网络权重初始化，实际上 pytorch 在构建网络时会自动初始化权重
def _initialize_weights(self):
for m in self.modules(): #遍历self.modules模块，通过self.modules模块会迭代定义的每一个层结构，判断其属于哪个类别
if isinstance(m, nn.Conv2d): #如果是卷积层，就用kaiming_normal对卷积权重w进行初始化
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
if m.bias is not None: #如果偏置不为空，就用0对其进行初始化
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.Linear): #如果是全连接层，那么采用normal对权重进行赋值，对偏置初始化为0
nn.init.normal_(m.weight, 0, 0.01)
nn.init.constant_(m.bias, 0)
2. train.py
与pytorch官方demo实现图像分类的train脚本差不多
数据预处理
data_transform = {
"train": transforms.Compose([transforms.RandomResizedCrop(224), #随机裁剪成224x224大小
transforms.RandomHorizontalFlip(), #随机翻转，水平方向
transforms.ToTensor(), #转化成Tensor
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))])} #标准化处理
载入数据集
#os.getcwd()获取当前文件所在的目录；os.path.join()将输入的路径连接在一起；..表示返回上层目录；../..表示返回上上层目录
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) #通过len()函数打印训练集有多少张图片
这里的数据集是在模型文件夹之外存放的，如果数据集存放在别的地方可以将data_root隐去，设置自己的路径
查看数据集
#查看数据集，查看之前要把batch_size修改成4，随机打乱shuffle改成True
# 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))
完整train.py代码
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") #有GPU就用，没有就不用
print("using {} device.".format(device))
#对数据进行预处理
data_transform = {
"train": transforms.Compose([transforms.RandomResizedCrop(224), #随机裁剪成224x224大小
transforms.RandomHorizontalFlip(), #随机翻转，水平方向
transforms.ToTensor(), #转化成Tensor
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))])} #标准化处理
#os.getcwd()获取当前文件所在的目录；os.path.join()将输入的路径连接在一起；..表示返回上层目录；../..表示返回上上层目录
data_root = os.path.abspath(os.path.join(os.getcwd(), "../..")) # get data root path获取数据集所在的根目录
image_path = os.path.join(data_root, "./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) #通过len()函数打印训练集有多少张图片
# # 字典，类别：索引 {'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())
# write dict into json file
json_str = json.dumps(cla_dict, indent=4) # 将 cla_dict 写入 json 文件中，方便在预测时读取信息
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_dataset
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修改成4，随机打乱shuffle改成True
# 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) #Adam优化器
epochs = 10
save_path = './AlexNet.pth'
best_acc = 0.0 #最佳准确率，在训练网络过程中保存准确率最高的模型
train_steps = len(train_loader)
for epoch in range(epochs):
# train 在网络搭建中使用了dropout方法，会在正向传播过程中随机失活一部分神经。但是只希望在训练过程中使用，在预测过程中不使用，因此通过net.train和net.eval来管理dropout(也可以管理BN层)
net.train() #在训练过程中调用net.train会启用dropout
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() #在验证过程中调用net.eval会关闭dropout
acc = 0.0 # accumulate accurate number / epoch
with torch.no_grad(): #禁止pytorch对参数跟踪，即在验证过程中不进行损失计算
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')
3. predict.py
def main():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#对载入图片进行预处理
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_path = "../tulip.jpg"
assert os.path.exists(img_path), "file: '{}' dose not exist.".format(img_path)
img = Image.open(img_path)
plt.imshow(img)
# [N, C, H, W]
img = data_transform(img)
# expand batch dimension
img = torch.unsqueeze(img, dim=0) #对图片扩充一个维度(0 batch)
# read class_indict
json_path = './class_indices.json'
assert os.path.exists(json_path), "file: '{}' dose not exist.".format(json_path)
json_file = open(json_path, "r")
class_indict = json.load(json_file)
# create model
model = AlexNet(num_classes=5).to(device)
# load model weights
weights_path = "./AlexNet.pth"
assert os.path.exists(weights_path), "file: '{}' dose not exist.".format(weights_path)
model.load_state_dict(torch.load(weights_path))
model.eval()
with torch.no_grad():
# predict class
output = torch.squeeze(model(img.to(device))).cpu()
predict = torch.softmax(output, dim=0)
predict_cla = torch.argmax(predict).numpy()
print_res = "class: {} prob: {:.3}".format(class_indict[str(predict_cla)],
predict[predict_cla].numpy())
plt.title(print_res)
for i in range(len(predict)):
print("class: {:10} prob: {:.3}".format(class_indict[str(i)],
predict[i].numpy()))
plt.show()