本文利用手写数字数据集搭建一个简单的分类神经网络,代码有详细的注释,帮助理解整个搭建流程。
代码链接:https://gitee.com/ury-u/pytorch.git
(lab_log.txt有详细介绍每个文件夹的作用)
案例演示:
import torch
import torchvision.datasets as dataset
import torchvision.transforms as transforms
import torch.utils.data as data_utils
# 1、定义数据 -- 对数据的加载和采样,并对数据的预处理
train_data = dataset.MNIST(root="mnist", # 存放在mnist文件夹下,
train=True, # 首先加载数据集
transform=transforms.ToTensor(), # 将数据转换为tensor
download=True, # 将数据下载
)
test_data = dataset.MNIST(root="mnist",
train=False,
transform=transforms.ToTensor(),
download=False,
)
''' 当数据量非常大的时候,不能一次性的将数据放在网络中进行训练,因此,可以通过循环按批次从数据集中选取一小部分数据放在网络中进行训练,而每一部分的数据称为batchsize '''
# batchsize
train_loader = data_utils.DataLoader(dataset=train_data,
batch_size=32, # 批处理样本的大小,满足当前显存约束
shuffle=True, # 将数据shuffle之后,可以消除数据因排列带来的问题
)
test_loader = data_utils.DataLoader(dataset=test_data,
batch_size=32,
shuffle=True,
)
# 2、定义网络结构
class CNN(torch.nn.Module): # 定义卷积神经网络模型 -- 输入图片为 28 * 28
def __init__(self):
super(CNN,self).__init__()
self.conv = torch.nn.Sequential(
torch.nn.Conv2d(1,32,kernel_size=5,padding=2), # 创建了一个卷积层 --
# 输入通道数为1(灰度图像)并输出通道数为32。kernel_size=5指定了卷积核的大小为5x5。padding=2表示在输入的边缘周围添加了2个零填充,以保持输出特征图的大小与输入特征图相同。
torch.nn.BatchNorm2d(32), # 创建一个批归一化层,应用于有32个输出通道的特征图
torch.nn.ReLU(), # 创建一个ReLU(修正线性单元)激活函数层
torch.nn.MaxPool2d(2), # 创建一个最大池化层,采用2x2的池化窗口,用于减小特征图的空间维度,并提取出最显著的特征
)
self.fc = torch.nn.Linear(14 * 14 * 32, 10) # 创建一个线性变换层,将输入的特征向量映射到输出的特征向量
# 在池化层之后,特征图的空间维度被减小为 14x14,并且有 32 个通道(或特征图)。因此,输入特征向量的维度为 14 * 14 * 32, 10 指定了输出特征向量的维度。
def forward(self,x): # 定义一个前向传播函数,用于执行神经网络的前向计算过程
out = self.conv(x) # 将输入 x 传递给之前定义的 self.conv 层(卷积层、批归一化层、ReLU激活函数层和最大池化层的组合)
out = out.view(out.size()[0],-1) # 将 out 转换为一个二维张量,其中每一行表示一个样本的特征向量
out = self.fc(out) # 将变形后的 out 输入到全连接层 self.fc 中进行线性变换,得到最终的输出
return out
cnn = CNN() # 初始化类
cnn = cnn.cuda() #将网络从CPU转到GPU
# 3、定义loss
loss_func = torch.nn.CrossEntropyLoss() # 由于是分类问题,因此定义交叉熵损失
# 4、优化函数
optimizer = torch.optim.Adam(cnn.parameters(),lr=0.01)
# 5、训练过程
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
for epoch in range(10):
for i, (images,labels) in enumerate(train_loader): # 通过 train_loader 迭代器遍历训练数据集,images 是输入图像的批次,labels 是对应的真实标签
images = images.to(device)
labels = labels.to(device)
outputs = cnn(images) # 将输入图像 images 通过卷积神经网络模型 cnn 进行前向传播,得到模型的输出结果
loss = loss_func(outputs,labels) # 计算模型outputs和真实标签之间的损失值 loss
# 反向传播、更新模型参数
optimizer.zero_grad()
loss.backward()
optimizer.step()
print("epoch is {}, ite is {}/{}, loss is {}".format(epoch+1, i,
len(train_data) // 32, # 总迭代次数
loss.item()))
# 6、eval/test
loss_test = 0
accuracy = 0
for i, (images,labels) in enumerate(test_loader):
images = images.to(device)
labels = labels.to(device)
outputs = cnn(images)
# [batchsize]
# outputs = batchsize * cls_num(10)
loss_test += loss_func(outputs,labels) # 累加每个测试样本的损失值
_, pred = outputs.max(1) # 取每个样本预测的最大概率对应的类别索引,并将其保存在 pred 中
accuracy += (pred == labels).sum().item() # 计算正确分类的样本数量,并将其累加到 accuracy 变量中
accuracy = accuracy / len(test_data)
loss_test = loss_test / (len(test_data) // 32) # 除以batchsize的总数量
print("epoch is {}, accuracy is {}, loss_test is {}".format(epoch + 1,
accuracy,
loss_test.item()))
torch.cuda.empty_cache()
# 7、save
torch.save(cnn,"model/mnist_model.pkl")
重新定义一个脚本,用于加载刚刚保存的模型。代码示例:
import torch
import torchvision.datasets as dataset
import torchvision.transforms as transforms
import torch.utils.data as data_utils
class CNN(torch.nn.Module):
def __init__(self):
super(CNN,self).__init__()
self.conv = torch.nn.Sequential(
torch.nn.Conv2d(1,32,kernel_size=5,padding=2),
torch.nn.BatchNorm2d(32),
torch.nn.ReLU(),
torch.nn.MaxPool2d(2),
)
self.fc = torch.nn.Linear(14 * 14 * 32, 10)
def forward(self,x):
out = self.conv(x)
out = out.view(out.size()[0],-1)
out = self.fc(out)
return out
test_data = dataset.MNIST(root="mnist",
train=False,
transform=transforms.ToTensor(),
download=False,
)
test_loader = data_utils.DataLoader(dataset=test_data,
batch_size=32,
shuffle=True,
)
cnn = torch.load("model/mnist_model.pkl") # 加载模型
cnn = cnn.cuda()
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
loss_test = 0
accuracy = 0
for i, (images,labels) in enumerate(test_loader):
images = images.to(device)
labels = labels.to(device)
outputs = cnn(images)
_, pred = outputs.max(1)
accuracy += (pred == labels).sum().item()
accuracy = accuracy / len(test_data)
print(accuracy)
但是在大多数情况下,我们一般会卷积神经网络的定义单独定义成一个CNN.py类,这样在使用的时候,我们直接导入该类。接下来我们将图片数据和预测结果都显示出来,代码示例:
import torch
import torchvision.datasets as dataset
import torchvision.transforms as transforms
import torch.utils.data as data_utils
from CNN import CNN # 导入CNN类
import cv2
test_data = dataset.MNIST(root="mnist",
train=False,
transform=transforms.ToTensor(),
download=False,
)
test_loader = data_utils.DataLoader(dataset=test_data,
batch_size=32,
shuffle=True,
)
cnn = torch.load("model/mnist_model.pkl") # 加载模型
cnn = cnn.cuda()
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
loss_test = 0
accuracy = 0
for i, (images,labels) in enumerate(test_loader):
images = images.to(device)
labels = labels.to(device)
outputs = cnn(images)
_, pred = outputs.max(1)
accuracy += (pred == labels).sum().item()
images = images.cpu().numpy() # 显示图片,从GPU转到CPU
labels = labels.cpu().numpy()
pred = pred.cpu().numpy()
# 那么现在就拿到了 batchsize * 1 * 28 * 28 维度的图片数据
for idx in range(images.shape[0]): # 对图片数据进行循环
im_data = images[idx] # 每次取一张图片
im_lable = labels[idx] # 并获取lable
im_pred = pred[idx]
im_data = im_data.transpose(1,2,0) # 将图片数据的维度顺序进行转置,从原始的通道-高度-宽度的顺序转换为高度-宽度-通道(HWC)的顺序
print("im_lable",im_lable)
print("pred",im_pred)
cv2.imshow("im_data",im_data) # 使用OpenCV库的 imshow 函数显示图片数据 im_data
cv2.waitKey(0)
accuracy = accuracy / len(test_data)
print(accuracy)