Pytorch：搭建卷积神经网络完成MNIST分类任务：
2023.7.18
MNIST百科：
MNIST数据集简介与使用_bwqiang的博客-CSDN博客
数据集官网：MNIST handwritten digit database, Yann LeCun, Corinna Cortes and Chris Burges
MNIST数据集获取并转换成图片格式：
数据集将按以图片和文件夹名为标签的形式保存：
代码：下载mnist数据集并转还为图片
import os
from PIL import Image
from torchvision import datasets, transforms
# 定义数据转换
transform = transforms.Compose([
transforms.ToTensor(), # 转换为张量
transforms.Normalize((0.5,), (0.5,)) # 标准化
])
# 下载并加载训练集和测试集
train_dataset = datasets.MNIST(root=os.getcwd(), train=True, transform=transform, download=True)
test_dataset = datasets.MNIST(root=os.getcwd(), train=False, transform=transform, download=True)
# 路径
train_path = './images/train'
test_path = './images/test'
# 将训练集中的图像保存为图片
for i in range(10):
file_name = train_path + os.sep + str(i)
if not os.path.exists(file_name):
os.mkdir(file_name)
for i in range(10):
file_name = test_path + os.sep + str(i)
if not os.path.exists(file_name):
os.mkdir(file_name)
for i, (image, label) in enumerate(train_dataset):
train_label = label
image_path = f'images/train/{train_label}/{i}.png'
image = image.squeeze().numpy() # 去除通道维度，并转换为 numpy 数组
image = (image * 0.5) + 0.5 # 反标准化，将范围调整为 [0, 1]
image = (image * 255).astype('uint8') # 将范围调整为 [0, 255]，并转换为整数类型
Image.fromarray(image).save(image_path)
# 将测试集中的图像保存为图片
for i, (image, label) in enumerate(test_dataset):
text_label = label
image_path = f'images/test/{text_label}/{i}.png'
image = image.squeeze().numpy() # 去除通道维度，并转换为 numpy 数组
image = (image * 0.5) + 0.5 # 反标准化，将范围调整为 [0, 1]
image = (image * 255).astype('uint8') # 将范围调整为 [0, 255]，并转换为整数类型
Image.fromarray(image).save(image_path)训练代码：
import os
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import torchvision.transforms as transforms
from PIL import Image
# 调动显卡进行计算
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
class MyDataset(torch.utils.data.Dataset):
def __init__(self, root_dir, transform=None):
self.root_dir = root_dir
self.transform = transform
self.names_list = []
for dirs in os.listdir(self.root_dir):
dir_path = self.root_dir + '/' + dirs
for imgs in os.listdir(dir_path):
img_path = dir_path + '/' + imgs
self.names_list.append((img_path, dirs))
def __len__(self):
return len(self.names_list)
def __getitem__(self, index):
image_path, label = self.names_list[index]
if not os.path.isfile(image_path):
print(image_path + '不存在该路径')
return None
image = Image.open(image_path)
label = np.array(label).astype(int)
label = torch.from_numpy(label)
if self.transform:
image = self.transform(image)
return image, label
# 定义卷积神经网络模型
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv1 = nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1)
self.relu = nn.ReLU()
self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
self.fc = nn.Linear(16 * 14 * 14, 10)
def forward(self, x):
x = self.conv1(x) # 卷积
x = self.relu(x) # 激活函数
x = self.maxpool(x) # 最大值池化
x = x.view(x.size(0), -1)
x = self.fc(x) # 全连接层
return x
# 加载手写数字数据集
train_dataset = MyDataset('./dataset/images/train', transform=transforms.ToTensor())
val_dataset = MyDataset('./dataset/images/val', transform=transforms.ToTensor())
# 定义超参数
batch_size = 8192 # 批处理大小
learning_rate = 0.001 # 学习率
num_epochs = 30 # 迭代次数
# 创建数据加载器
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
val_loader = torch.utils.data.DataLoader(dataset=val_dataset, batch_size=batch_size, shuffle=False)
# 实例化模型、损失函数和优化器
model = CNN().to(device)
criterion = nn.CrossEntropyLoss() # 损失函数
optimizer = optim.Adam(model.parameters(), lr=learning_rate) # 优化器
# 记录验证的次数
total_train_step = 0
total_val_step = 0
# 模型训练和验证
print("-------------TRAINING-------------")
total_step = len(train_loader)
for epoch in range(num_epochs):
print("Epoch=", epoch)
for i, (images, labels) in enumerate(train_loader):
images = images.to(device)
labels = labels.to(device)
output = model(images)
loss = criterion(output, labels.long())
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_train_step = total_train_step + 1
print("train_times:{},Loss:{}".format(total_train_step, loss.item()))
# 测试验证
total_val_loss = 0
total_accuracy = 0
with torch.no_grad():
for i, (images, labels) in enumerate(val_loader):
images = images.to(device)
labels = labels.to(device)
outputs = model(images)
loss = criterion(outputs, labels.long())
total_val_loss = total_val_loss + loss.item() # 计算损失值的和
accuracy = 0
for j in labels: # 计算精确度的和
if outputs.argmax(1)[j] == labels[j]:
accuracy = accuracy + 1
total_accuracy = total_accuracy + accuracy
print('Accuracy =', float(total_accuracy / len(val_dataset))) # 输出正确率
torch.save(model, "cnn_{}.pth".format(epoch)) # 模型保存
# # 模型评估
# with torch.no_grad():
# correct = 0
# total = 0
# for images, labels in test_loader:
# outputs = model(images)
# _, predicted = torch.max(outputs.data, 1)
# total += labels.size(0)
# correct += (predicted == labels).sum().item()
测试代码：
import torch
from torchvision import transforms
import torch.nn as nn
import os
from PIL import Image
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') # 判断是否有GPU
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv1 = nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1)
self.relu = nn.ReLU()
self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
self.fc = nn.Linear(16 * 14 * 14, 10)
def forward(self, x):
x = self.conv1(x) # 卷积
x = self.relu(x) # 激活函数
x = self.maxpool(x) # 最大值池化
x = x.view(x.size(0), -1)
x = self.fc(x) # 全连接层
return x
model = torch.load('cnn.pth') # 加载模型
path = "./dataset/images/test/" # 测试集
imgs = os.listdir(path)
test_num = len(imgs)
print(f"test_dataset_quantity={test_num}")
for img_name in imgs:
img = Image.open(path + img_name)
test_transform = transforms.Compose([transforms.ToTensor()])
img = test_transform(img)
img = img.to(device)
img = img.unsqueeze(0)
outputs = model(img) # 将图片输入到模型中
_, predicted = outputs.max(1)
pred_type = predicted.item()
print(img_name, 'pred_type:', pred_type)分类正确率不错：
