AlexNet基于MNIST数据集的代码实现

鉴于原论文中使用的数据集过于庞大&＃xff0c;分类过多&＃xff0c;目前手头的设备运行是在过于缓慢&＃xff0c;折中考虑尝试使用MNIST的数据集实现AlexNet

import torch, torchvision
import torchvision.transforms as transforms
import torch.nn as nn
from torch import optim
import matplotlib.pyplot as plt
%matplotlib inline
import copy


# 超参数设置
EPOCH &＃61; 10
BATCH_SIZE &＃61; 64
LR &＃61; 0.01


transform &＃61; transforms.ToTensor()


数据集

通过torchvision下载数据集

trainset &＃61; torchvision.datasets.MNIST(root&＃61;&＃39;../data&＃39;, train&＃61;True, download&＃61;True, transform&＃61;transform)testset &＃61; torchvision.datasets.MNIST(root&＃61;&＃39;../data&＃39;, train&＃61;True, transform&＃61;transform)


C:\Users\Administrator\AppData\Roaming\Python\Python36\site-packages\torchvision\datasets\mnist.py:498: UserWarning: The given NumPy array is not writeable, and PyTorch does not support non-writeable tensors. This means you can write to the underlying (supposedly non-writeable) NumPy array using the tensor. You may want to copy the array to protect its data or make it writeable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at ..\torch\csrc\utils\tensor_numpy.cpp:180.)return torch.from_numpy(parsed.astype(m[2], copy&＃61;False)).view(*s)


trainloader &＃61; torch.utils.data.DataLoader(trainset, batch_size&＃61;BATCH_SIZE, shuffle&＃61;True)testloader &＃61; torch.utils.data.DataLoader(testset, batch_size&＃61;BATCH_SIZE, shuffle&＃61;False)


绘图查看

plt.imshow(trainset[4][0][0], cmap&＃61;&＃39;gray&＃39;)


查看数据格式

trainset[0][0].shape


torch.Size([1, 28, 28])


为了通用&＃xff0c;设置一个device&＃xff0c;如果有显卡并配置好了cuda环境&＃xff0c;那么就选择为cuda&＃xff0c;否则为cpu

device &＃61; torch.device(&＃39;cuda&＃39; if torch.cuda.is_available() else &＃39;cpu&＃39;)
device


device(type&＃61;&＃39;cpu&＃39;)


AlexNet

同样地&＃xff0c;仿照AlexNet&＃xff0c;设置了五个卷积层和三个全连接层构建一个深度卷积神经网络&＃xff0c;网络的定义是重写nn.Module实现的&＃xff0c;卷积层和全连接层之间将数据通过view拉平^[1]

class AlexNet(nn.Module):def __init__(self,width_mult&＃61;1):super(AlexNet, self).__init__()self.layer1 &＃61; nn.Sequential(nn.Conv2d(1, 32, kernel_size&＃61;3, padding&＃61;1), # 32*28*28nn.MaxPool2d(kernel_size&＃61;2, stride&＃61;2), # 32*14*14nn.ReLU(inplace&＃61;True),)self.layer2 &＃61; nn.Sequential(nn.Conv2d(32, 64, kernel_size&＃61;3, padding&＃61;1), # 64*14*14nn.MaxPool2d(kernel_size&＃61;2, stride&＃61;2), # 64*7*7nn.ReLU(inplace&＃61;True),)self.layer3 &＃61; nn.Sequential(nn.Conv2d(64, 128, kernel_size&＃61;3, padding&＃61;1), # 128*7*7)self.layer4 &＃61; nn.Sequential(nn.Conv2d(128, 256, kernel_size&＃61;3, padding&＃61;1), # 256*7*7)self.layer5 &＃61; nn.Sequential(nn.Conv2d(256, 256, kernel_size&＃61;3, padding&＃61;1), # 256*7*7nn.MaxPool2d(kernel_size&＃61;3, stride&＃61;2), # 256*3*3nn.ReLU(inplace&＃61;True),)self.fc1 &＃61; nn.Linear(256*3*3, 1024)self.fc2 &＃61; nn.Linear(1024, 512)self.fc3 &＃61; nn.Linear(512, 10)def forward(self, x):x &＃61; self.layer1(x)x &＃61; self.layer2(x)x &＃61; self.layer3(x)x &＃61; self.layer4(x)x &＃61; self.layer5(x)x &＃61; x.view(-1, 256*3*3)x &＃61; self.fc1(x)x &＃61; self.fc2(x)x &＃61; self.fc3(x)return x


设置超参数

EPOCH &＃61; 5
BATCH_SIZE &＃61; 128
LR &＃61; 0.01


def validate(model, data):total &＃61; 0correct &＃61; 0for i, (images, labels) in enumerate(data):images &＃61; images.to(device)x &＃61; net(images)value, pred &＃61; torch.max(x,1)pred &＃61; pred.data.cpu()total &＃43;&＃61; x.size(0)correct &＃43;&＃61; torch.sum(pred &＃61;&＃61; labels)return correct*100./total


初始化模型并将模型放到device上&＃xff0c;如果有显卡就在cuda上&＃xff0c;如果没有&＃xff0c;那么在cpu

如果是纯cpu训练&＃xff0c;速度十分感人

net &＃61; AlexNet().to(device)


# alexnet训练
def train():# 定义损失函数为交叉熵损失&＃xff0c;优化方法为SGDcriterion &＃61; nn.CrossEntropyLoss() optimizer &＃61; optim.SGD(net.parameters(), lr&＃61;LR, momentum&＃61;0.9)max_accuracy&＃61;0accuracies&＃61;[]for epoch in range(EPOCH):for i, (images,labels) in enumerate(trainloader):images &＃61; images.to(device)labels &＃61; labels.to(device)optimizer.zero_grad()outputs &＃61; net(images)loss &＃61; criterion(outputs, labels)loss_item &＃61; loss.item()loss.backward()optimizer.step()accuracy &＃61; float(validate(criterion, testloader))accuracies.append(accuracy)print("Epoch %d accuracy: %f loss: %f" % (epoch, accuracy, loss_item))if accuracy > max_accuracy:best_model &＃61; copy.deepcopy(criterion)max_accuracy &＃61; accuracyprint("Saving Best Model with Accuracy: ", accuracy)print(&＃39;Epoch:&＃39;, epoch&＃43;1, "Accuracy :", accuracy, &＃39;%&＃39;)plt.plot(accuracies)return best_model


这一行代码是调用之前的train函数训练神经网络&＃xff0c;初始化设置的epoch是5&＃xff0c;大概也可以训练一个准确度较高的模型

alexnet &＃61; train()


为了防止断点或者bug导致jupyter重启之后重新训练模型&＃xff0c;这一点经常遇到&＃xff0c;本代码是在google的colab上训练的&＃xff0c;为了保存训练的结果&＃xff0c;还是将模型保存为pkl文件&＃xff0c;这样本地就不用训练&＃xff0c;直接调用训练之后的模型&＃xff0c;之前尝试直接保存整个模型&＃xff0c;但是会有莫名其妙的bug&＃xff0c;暂时没有解决。这里尝试了另一种保存模型的方式^[2]&＃xff0c;直接保存模型的参数&＃xff0c;然后将参数传递到初始化的模型架构中&＃xff0c;如下所示&＃xff1a;

# 保存模型参数
torch.save(alexnet, &＃39;../models/alexnet.pkl&＃39;)


# 加载模型
alexnet &＃61; AlexNet()
alexnet.load_state_dict(torch.load(&＃39;../models/alexnet.pkl&＃39;))


AlexNet((layer1): Sequential((0): Conv2d(1, 32, kernel_size&＃61;(3, 3), stride&＃61;(1, 1), padding&＃61;(1, 1))(1): MaxPool2d(kernel_size&＃61;2, stride&＃61;2, padding&＃61;0, dilation&＃61;1, ceil_mode&＃61;False)(2): ReLU(inplace&＃61;True))(layer2): Sequential((0): Conv2d(32, 64, kernel_size&＃61;(3, 3), stride&＃61;(1, 1), padding&＃61;(1, 1))(1): MaxPool2d(kernel_size&＃61;2, stride&＃61;2, padding&＃61;0, dilation&＃61;1, ceil_mode&＃61;False)(2): ReLU(inplace&＃61;True))(layer3): Sequential((0): Conv2d(64, 128, kernel_size&＃61;(3, 3), stride&＃61;(1, 1), padding&＃61;(1, 1)))(layer4): Sequential((0): Conv2d(128, 256, kernel_size&＃61;(3, 3), stride&＃61;(1, 1), padding&＃61;(1, 1)))(layer5): Sequential((0): Conv2d(256, 256, kernel_size&＃61;(3, 3), stride&＃61;(1, 1), padding&＃61;(1, 1))(1): MaxPool2d(kernel_size&＃61;3, stride&＃61;2, padding&＃61;0, dilation&＃61;1, ceil_mode&＃61;False)(2): ReLU(inplace&＃61;True))(fc1): Linear(in_features&＃61;2304, out_features&＃61;1024, bias&＃61;True)(fc2): Linear(in_features&＃61;1024, out_features&＃61;512, bias&＃61;True)(fc3): Linear(in_features&＃61;512, out_features&＃61;10, bias&＃61;True)
)


为直观的查看效果&＃xff0c;选择一组测试集图片查看分类效果

plt.figure(figsize&＃61;(14, 14))
for i, (image, label) in enumerate(testloader):predict &＃61; torch.argmax(alexnet(image), axis&＃61;1)print((predict &＃61;&＃61; label).sum()/label.shape[0])for j in range(image.shape[0]):plt.subplot(8, 8, j&＃43;1)plt.imshow(image[j, 0], cmap&＃61;&＃39;gray&＃39;)plt.title(predict[j].item())plt.axis(&＃39;off&＃39;)if i&＃61;&＃61;1:break


tensor(1.)


参考文献

• [1] Sowndharya206/alexnet

• [2] SAVE AND LOAD THE MODEL