PyTorch复现AlexNet学习笔记

PyTorch复现AlexNet学习笔记

一篇简单的学习笔记，实现五类花分类

这里只介绍复现的工作，如果想了解更多有关网络的细节，请去看论文《ImageNet Classification with Deep Convolutional Neural Networks》

简单说明下数据集，下载链接

下载解压数据集

                  

一、环境准备

可以去看上一篇博客，里面写的很详细了，并且推荐了一篇炮哥的环境搭建环境

二、模型搭建、训练

1.整体框图

AlexNet整体框图，padding，stride，需要根据论文计算，前后卷积大小没变，一般padding=2

2.net.py

网络整体结构代码

 1 import torch
 2 from torch import nn
 3 import torch.nn.functional as F
 4
 5 class MyAlexNet(nn.Module):
 6     def __init__(self,num_classes):
 7         super(MyAlexNet, self).__init__()
 8         self.c1 = nn.Conv2d(in_channels=3,out_channels=48,kernel_size=11,stride=4,padding=2)
 9         self.ReLu = nn.ReLU()
10         self.c2 = nn.Conv2d(in_channels=48,out_channels=128,kernel_size=5,stride=1,padding=2)
11         self.s2 = nn.MaxPool2d(2)
12         self.c3 = nn.Conv2d(in_channels=128,out_channels=192,kernel_size=3,stride=1,padding=2)
13         self.s3 = nn.MaxPool2d(2)
14         self.c4 = nn.Conv2d(in_channels=192,out_channels=192,kernel_size=3,stride=1,padding=1)
15         self.c5 = nn.Conv2d(in_channels=192,out_channels=128,kernel_size=3,stride=1,padding=1)
16         self.s5 = nn.MaxPool2d(kernel_size=3,stride=2)
17         self.flatten = nn.Flatten()
18         self.f6 = nn.Linear(4608,2048)#经过池化后的神经元个数(13-3)/2+1=6,6*6*128=4608
19         self.f7 = nn.Linear(2048,2048)
20         self.f8 = nn.Linear(2048,1000)
21         self.f9 = nn.Linear(1000,num_classes)#分类类别数
22
23     def forward(self,x):
24         x = self.ReLu(self.c1(x))
25         x = self.ReLu(self.c2(x))
26         x = self.s2(x)
27         x = self.ReLu(self.c3(x))
28         x = self.s3(x)
29         x = self.ReLu(self.c4(x))
30         x = self.ReLu(self.c5(x))
31         x = self.s5(x)
32         x = self.flatten(x)
33         x = self.f6(x)
34         x = F.dropout(x,0.5)
35         x = self.f7(x)
36         x = F.dropout(x,0.5)
37         x = self.f8(x)
38         x = F.dropout(x,0.5)
39         x = self.f9(x)
40
41         return x
42
43 if __name__ =="__main__":
44     x = torch.rand([1, 3, 224, 224])
45     model = MyAlexNet(num_classes=5)
46     y = model(x)
47     print(y)
48     # 统计模型参数  total param num 16632442
49     # sum = 0
50     # for name, param in model.named_parameters():
51     #     num = 1
52     #     for size in param.shape:
53     #         num *= size
54     #     sum += num
55     #     # print("{:30s} : {}".format(name, param.shape))
56     # print("total param num {}".format(sum))  # total param num 134,281,029

net.py

写完后保存，运行可以检查是否报错

3.数据划分

分好后的数据集

运行下面代码将数据按一定比例，划分为训练集和验证集

 1 import os
 2 from shutil import copy
 3 import random
 4
 5
 6 def mkfile(file):
 7     if not os.path.exists(file):
 8         os.makedirs(file)
 9
10
11 # 获取data文件夹下所有文件夹名（即需要分类的类名）
12 file_path = 'data'  #需要划分数据集的路径
13 flower_class = [cla for cla in os.listdir(file_path)]
14
15 # 创建 训练集train 文件夹，并由类名在其目录下创建5个子目录
16 mkfile('data/train')
17 for cla in flower_class:
18     mkfile('data/train/' + cla)
19
20 # 创建 验证集val 文件夹，并由类名在其目录下创建子目录
21 mkfile('data/val')
22 for cla in flower_class:
23     mkfile('data/val/' + cla)
24
25 # 划分比例
26 split_rate = 0.2 #20%为验证集
27
28 # 遍历所有类别的全部图像并按比例分成训练集和验证集
29 for cla in flower_class:
30     cla_path = file_path + '/' + cla + '/'  # 某一类别的子目录
31     images = os.listdir(cla_path)  # iamges 列表存储了该目录下所有图像的名称
32     num = len(images)
33     eval_index = random.sample(images, k=int(num * split_rate))  # 从images列表中随机抽取 k 个图像名称
34     for index, image in enumerate(images):
35         # eval_index 中保存验证集val的图像名称
36         if image in eval_index:
37             image_path = cla_path + image
38             new_path = 'data/val/' + cla
39             copy(image_path, new_path)  # 将选中的图像复制到新路径
40
41         # 其余的图像保存在训练集train中
42         else:
43             image_path = cla_path + image
44             new_path = 'data/train/' + cla
45             copy(image_path, new_path)
46         print("\r[{}] processing [{}/{}]".format(cla, index + 1, num), end="")  # processing bar
47     print()
48
49 print("processing done!")

数据划分的代码

4.train.py

训练的代码，训练结束后画出训练集和验证集的loss,准确度，60轮，batch-size=16，SGD优化算法，学习率0.01，10轮变为原来的0.5。

  1 #修改后加进度条的代码
  2 import json
  3 import torch
  4 from torch import nn
  5 from NET import MyAlexNet
  6 import numpy as np
  7
  8 from tqdm import tqdm#用于画进度条
  9
 10 from torch.optim import lr_scheduler
 11
 12 import os
 13 import sys
 14
 15 from torchvision import transforms
 16 from torchvision.datasets import ImageFolder
 17 from torch.utils.data import DataLoader
 18
 19 import matplotlib.pyplot as plt
 20
 21 # 解决中文显示问题
 22 plt.rcParams['font.sans-serif'] = ['SimHei']
 23 plt.rcParams['axes.unicode_minus'] = False
 24
 25 # 如果显卡可用，则用显卡进行训练
 26 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 27 print("using {} device".format(device))
 28
 29
 30 # 将图像RGB三个通道的像素值分别减去0.5,再除以0.5.从而将所有的像素值固定在[-1,1]范围内
 31 #normalize = transforms.Normalize(std=[0.5,0.5,0.5],mean=[0.5,0.5,0.5])#image=(image-mean)/std
 32 data_transform = {
 33     "train":transforms.Compose([
 34             transforms.Resize((224,224)),#裁剪为224*224
 35             transforms.RandomVerticalFlip(),#随机垂直旋转
 36             transforms.ToTensor(),#将0-255范围内的像素转为0-1范围内的tensor
 37             transforms.Normalize(std=[0.5,0.5,0.5],mean=[0.5,0.5,0.5])#归一化
 38 ]),
 39     "val":transforms.Compose([
 40           transforms.Resize((224,224)),#裁剪为224*224
 41           transforms.ToTensor(),#将0-255范围内的像素转为0-1范围内的tensor
 42           transforms.Normalize(std=[0.5,0.5,0.5],mean=[0.5,0.5,0.5])#归一化
 43 ])}
 44
 45 #数据集路径
 46 ROOT_TRAIN = 'data/train'
 47 ROOT_TEST = 'data/val'
 48
 49 batch_size = 16
 50
 51 train_dataset = ImageFolder(ROOT_TRAIN,transform=data_transform["train"])#ImageFolder()根据文件夹名来对图像添加标签
 52 val_dataset = ImageFolder(ROOT_TEST,transform=data_transform["val"])#可以利用print(val_dataset.imgs)对象查看,返回列表形式('data/val\\cat\\110.jpg', 0)
 53 #print(val_dataset.imgs)
 54
 55 # nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])  # number of workers
 56 # print('Using {} dataloader workers every process'.format(nw))
 57
 58 train_dataloader = DataLoader(train_dataset,batch_size=batch_size,shuffle=True)
 59 val_dataloader = DataLoader(val_dataset,batch_size=batch_size,shuffle=True)
 60
 61 flow_list = train_dataset.class_to_idx#转换维字典,train_dataset里有这个对象
 62 # {'daisy':0, 'dandelion':1, 'roses':2, 'sunflower':3, 'tulips':4}
 63 cla_dict = dict((val,key) for key,val in flow_list.items())#键值对转换
 64 #{0: 'daisy', 1: 'dandelion', 2: 'roses', 3: 'sunflowers', 4: 'tulips'}
 65 # write dict into json file
 66 json_str = json.dumps(cla_dict, indent=4)
 67 with open('class_indices.json', 'w') as json_file:
 68     json_file.write(json_str)  # 保存json文件(好处，方便转换为其它类型数据)用于预测用
 69
 70 train_num = len(train_dataset)
 71 val_num = len(val_dataset)
 72 print("using {} images for training, {} images for validation.".format(train_num,val_num))
 73
 74 # 调用net里面的定义的网络模型， 如果GPU可用则将模型转到GPU
 75 model = MyAlexNet(num_classes=5).to(device)
 76
 77 #加载预训练模型
 78 # weights_path = "save_model/best_model.pth"
 79 # assert os.path.exists(weights_path), "file: '{}' dose not exist.".format(weights_path)
 80 # missing_keys, unexpected_keys = net.load_state_dict(torch.load(weights_path,),strict=False)
 81
 82 #定义损失函数
 83 loss_fn = nn.CrossEntropyLoss()
 84
 85 #定义优化器
 86 optimizer = torch.optim.SGD(model.parameters(),lr=0.01,momentum=0.9)#googlenet用的是adam
 87 # 学习率每隔10epoch变为原来的0.5
 88 lr_scheduler = lr_scheduler.StepLR(optimizer,step_size=10,gamma=0.5)
 89
 90 #定义训练函数
 91 def train(dataloader,model,loss_fn,optimizer,i,epoch):
 92     model.train()
 93     loss,current,n = 0.0,0.0,0
 94     train_bar = tqdm(dataloader,file=sys.stdout)#输出方式，默认为sys.stderr
 95     for batch,(x,y) in enumerate(train_bar):#enumerate()默认两个参数，第一个用于记录序号，默认0开始，第二个参数(x,y)才是需要遍历元素(dataloder)的值
 96         #前向传播
 97         image,y = x.to(device),y.to(device)
 98         output = model(image)
 99         cur_loss = loss_fn(output,y)
100         _,pred = torch.max(output,axis=-1)
101         cur_acc = torch.sum(y==pred)/output.shape[0]
102         #反向传播
103         optimizer.zero_grad()#梯度归零
104         cur_loss.backward()
105         optimizer.step()
106         loss += cur_loss
107         current += cur_acc
108         n += 1
109         train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(i + 1, epoch, cur_loss)
110     train_loss = loss / n
111     train_acc = current / n
112     print(f'train_loss:{train_loss}')
113     print(f'train_acc:{train_acc}')
114     return  train_loss,train_acc
115
116 #定义验证函数
117 def val(dataloader,model,loss_fn,i,epcho):
118     #转换为验证模型
119     model.eval()
120     loss, current, n = 0.0, 0.0, 0
121     with torch.no_grad():
122         val_bar = tqdm(dataloader,file=sys.stdout)
123         for batch, (x, y) in enumerate(val_bar):  # enumerate()默认两个参数，第一个用于记录序号，默认0开始，第二个参数(x,y)才是需要遍历元素(dataloder)的值
124             # 前向传播
125             image, y = x.to(device), y.to(device)
126             output = model(image)
127             cur_loss = loss_fn(output, y)
128             _, pred = torch.max(output, axis=-1)
129             cur_acc = torch.sum(y == pred) / output.shape[0]
130             loss += cur_loss
131             current += cur_acc
132             n += 1
133             val_bar.desc = "val epoch[{}/{}] loss:{:.3f}".format(i + 1, epoch, cur_loss)
134         val_loss = loss / n
135         val_acc = current / n
136         print(f'val_loss:{val_loss}')
137         print(f'val_acc:{val_acc}')
138         return val_loss, val_acc
139
140 #画图函数
141 def matplot_loss(train_loss,val_loss):
142     plt.figure()  # 声明一个新画布，这样两张图像的结果就不会出现重叠
143     plt.plot(train_loss,label='train_loss')#画图
144     plt.plot(val_loss, label='val_loss')
145     plt.legend(loc='best')#图例
146     plt.ylabel('loss',fontsize=12)
147     plt.xlabel('epoch',fontsize=12)
148     plt.title("训练集和验证集loss对比图")
149     plt.savefig('result/loss.jpg')
150
151 def matplot_acc(train_acc,val_acc):
152     plt.figure()  # 声明一个新画布，这样两张图像的结果就不会出现重叠
153     plt.plot(train_acc, label='train_acc')  # 画图
154     plt.plot(val_acc, label='val_acc')
155     plt.legend(loc='best')  # 图例
156     plt.ylabel('acc', fontsize=12)
157     plt.xlabel('epoch', fontsize=12)
158     plt.title("训练集和验证集acc对比图")
159     plt.savefig('result/acc.jpg')
160
161 #开始训练
162 train_loss_list = []
163 val_loss_list = []
164 train_acc_list = []
165 val_acc_list = []
166
167 epoch = 60
168 max_acc = 0
169
170 for i in range(epoch):
171     lr_scheduler.step()#学习率迭代，10epoch变为原来的0.5
172     train_loss,train_acc = train(train_dataloader,model,loss_fn,optimizer,i,epoch)
173     val_loss,val_acc = val(val_dataloader,model,loss_fn,i,epoch)
174
175     train_loss_list.append(train_loss)
176     train_acc_list.append(train_acc)
177     val_loss_list.append(val_loss)
178     val_acc_list.append(val_acc)
179     #保存最好的模型权重
180     if val_acc >max_acc:
181         folder = 'save_model'
182         if not os.path.exists(folder):
183             os.mkdir('save_model')
184         max_acc = val_acc
185         print(f'save best model,第{i+1}轮')
186         torch.save(model.state_dict(),'save_model/best_model.pth')#保存
187     #保存最后一轮
188     if i == epoch - 1:
189         torch.save(model.state_dict(), 'save_model/last_model.pth')  # 保存
190 print("done")
191
192 #画图
193 matplot_loss(train_loss_list,val_loss_list)
194 matplot_acc(train_acc_list,val_acc_list)

train.py

最后一轮的结果

训练结束后可以得到训练集和验证集的loss，acc对比图：

简单的评估下：模型在25轮左右，模型对训练集过拟合了。

如果想提高测试集准确度，需要去采用些手段来防止模型过拟合，比如正则化，数据增强等

三、模型测试

测试代码，这里用的测试集其实是之前训练验证集，本来是要另外创建一个的

 1 import os
 2 import json
 3 import torch
 4 from PIL import Image
 5 from torchvision import transforms
 6 import matplotlib.pyplot as plt
 7 from NET import MyAlexNet
 8
 9 def main():
10     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
11
12     data_transform = transforms.Compose([
13         transforms.Resize((224,224)),
14         transforms.ToTensor(),
15         transforms.Normalize((0.5, 0.5, 0.5),(0.5, 0.5, 0.5))
16     ])
17     #load image
18     img_path = "data/val/tulips/8677713853_1312f65e71.jpg"
19     assert os.path.exists(img_path),"file:'{}' dose not exist. ".format(img_path)
20     img = Image.open(img_path)
21     plt.imshow(img)
22
23     #[N, C, H, W]归一化
24     img = data_transform(img)
25     # expand batch dimension
26     img = torch.unsqueeze(img,dim=0)
27
28     # read class_indict
29     json_path = './class_indices.json'
30     assert os.path.exists(json_path), "file: '{}' dose not exist.".format(json_path)
31
32     with open(json_path,"r") as f:
33         class_indict = json.load(f)
34
35     #实例化模型
36     model = MyAlexNet(num_classes=5).to(device)
37
38     #加载权重
39     weights_path = "save_model/best_model.pth"
40     assert os.path.exists(weights_path), "file: '{}' dose not exist.".format(weights_path)
41     missing_keys,unexpected_keys = model.load_state_dict(torch.load(weights_path,map_location=device),
42                                                          strict=False)
43     model.eval()
44     with torch.no_grad():
45         #预测
46         output = torch.squeeze(model(img.to(device))).cpu()
47         predict = torch.softmax(output, dim=0)
48         predict_cla = torch.argmax(predict).numpy()
49     #最大概率结果
50     print_res = "class: {}   prob: {:.3}".format(class_indict[str(predict_cla)],
51                                                  predict[predict_cla].numpy())
52     #前10个类别
53     plt.title(print_res)
54     for i in range(len(predict)):
55         print("class: {:10}   prob: {:.3}".format(class_indict[str(i)],
56                                                   predict[i].numpy()))
57     plt.show()
58 if __name__=="__main__":
59     os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
60     main()

test.py

运行代码后，对模型进行推理，去网上找几张图片

下面是一张蒲公英照片，以及5类花预测的概率显示(右边)

总结

流程还是很顺利的，就是最后模型对训练集过拟合了，但精度还是很高的。

自己敲一下代码，会学到很多不懂的东西

比如ImageFolder()这个函数，是按照文件夹名字，来给文件夹里的数据打上标签

可以利用print(val_dataset.imgs)对象查看,返回列表形式('data/val\\cat\\110.jpg', 0)

最后，多看，多学，多试，总有一天你会称为大佬！