MNIST手写数字卷积识别
完整Demo代码
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0' #指定GPU编
import torch
import numpy as np
import unet
import matplotlib.pyplot as plt
from tqdm import tqdm # 进度条工具
batch_size = 320 #设定每次训练的批次数
epochs = 1024 #设定训练次数
device = "cpu" #Pytorch的特性,需要指定计算的硬件,如果没有GPU的存在,就使用CPU进行计算
#device = "cuda" #在这里读者默认使用GPU,如果读者出现运行问题可以将其改成cpu模式
model = unet.Unet() #导入Unet模型
model = model.to(device) #将计算模型传入GPU硬件等待计算
#model = torch.compile(model) #Pytorch2.0的特性,加速计算速度 选择使用内容
optimizer = torch.optim.Adam(model.parameters(), lr=2e-5) #设定优化函数
print("Finish1")
#载入数据
x_train = np.load("./data_book/x_train.npy")
y_train_label = np.load("./data_book/y_train_label.npy")
print("原始数据读取完成")
x_train_batch = []
for i in range(len(y_train_label)):
if y_train_label[i] <= 10: #为了加速演示作者只对数据集中的小于2的数字,也就是0和1进行运行,读者可以自行增加训练个数
x_train_batch.append(x_train[i])
x_train = np.reshape(x_train_batch, [-1, 1, 28, 28]) #修正数据输入维度:([30596, 28, 28])
#x_train /= 512.
x_train = x_train / 512.
train_length = len(x_train) * 20 #增加数据的单词循环次数
#state_dict = torch.load("./saver/unet.pth")
#model.load_state_dict(state_dict)
for epoch in range(2):
train_num = train_length // batch_size #计算有多少批次数 //表示整除
train_loss = 0 #用于损失函数的统计
for i in tqdm(range(train_num)): #开始循环训练
x_imgs_batch = [] #创建数据的临时存储位置
x_step_batch = []
y_batch = []
# 对每个批次内的数据进行处理
for b in range(batch_size):
img = x_train[np.random.randint(x_train.shape[0])] #提取单个图片内容
x = img
y = img
x_imgs_batch.append(x)
y_batch.append(y)
#将批次数据转化为Pytorch对应的tensor格式并将其传入GPU中
x_imgs_batch = torch.tensor(x_imgs_batch).float().to(device)
y_batch = torch.tensor(y_batch).float().to(device)
pred = model(x_imgs_batch) #对模型进行正向计算
loss = torch.nn.MSELoss(reduction="sum")(pred, y_batch)*100. #使用损失函数进行计算
#这里读者记住下面就是固定格式,一般而言这样使用即可
optimizer.zero_grad() #对结果进行优化计算
loss.backward() #损失值的反向传播
optimizer.step() #对参数进行更新
train_loss += loss.item() #记录每个批次的损失值
#计算并打印损失值
train_loss /= train_num
print("train_loss:", train_loss)
if epoch%6 == 0:
torch.save(model.state_dict(),"./saver/unet.pth")
#下面是对数据进行打印
image = x_train[np.random.randint(x_train.shape[0])] #随机挑选一条数据进行计算
image = np.reshape(image,[1,1,28,28]) #修正数据维度
image = torch.tensor(image).float().to(device) #挑选的数据传入硬件中等待计算
image = model(image) #使用模型对数据进行计算
image = torch.reshape(image, shape=[28,28]) #修正模型输出结果
image = image.detach().cpu().numpy() #将计算结果导入CPU中进行后续计算或者展示
#展示或计算数据结果
plt.imshow(image)
plt.savefig(f"./img/img_{epoch}.jpg")
在文件目录下创建unet.py 封装的模型类
import torch
import einops.layers.torch as elt
class Unet(torch.nn.Module):
def __init__(self):
super(Unet, self).__init__()
#模块化结构,这也是后面常用到的模型结构
self.first_block_down = torch.nn.Sequential(
torch.nn.Conv2d(in_channels=1,out_channels=32,kernel_size=3,padding=1),torch.nn.GELU(),
torch.nn.MaxPool2d(kernel_size=2,stride=2)
)
self.second_block_down = torch.nn.Sequential(
torch.nn.Conv2d(in_channels=32,out_channels=64,kernel_size=3,padding=1),torch.nn.GELU(),
torch.nn.MaxPool2d(kernel_size=2,stride=2)
)
self.latent_space_block = torch.nn.Sequential(
torch.nn.Conv2d(in_channels=64,out_channels=128,kernel_size=3,padding=1),torch.nn.GELU(),
)
self.second_block_up = torch.nn.Sequential(
torch.nn.Upsample(scale_factor=2),
torch.nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3, padding=1), torch.nn.GELU(),
)
self.first_block_up = torch.nn.Sequential(
torch.nn.Upsample(scale_factor=2),
torch.nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, padding=1), torch.nn.GELU(),
)
self.convUP_end = torch.nn.Sequential(
torch.nn.Conv2d(in_channels=32,out_channels=1,kernel_size=3,padding=1),
torch.nn.Tanh()
)
def forward(self,img_tensor):
image = img_tensor
image = self.first_block_down(image)#;print(image.shape) # torch.Size([5, 32, 14, 14])
image = self.second_block_down(image)#;print(image.shape) # torch.Size([5, 16, 7, 7])
image = self.latent_space_block(image)#;print(image.shape) # torch.Size([5, 8, 7, 7])
image = self.second_block_up(image)#;print(image.shape) # torch.Size([5, 16, 14, 14])
image = self.first_block_up(image)#;print(image.shape) # torch.Size([5, 32, 28, 28])
image = self.convUP_end(image)#;print(image.shape) # torch.Size([5, 32, 28, 28])
return image
if __name__ == '__main__':
image = torch.randn(size=(5,1,28,28))
unet_model = Unet()
torch.save(unet_model, './unet_model.pth')
原始数据读取完成
100%|██████████| 3750/3750 [1:37:04<00:00, 1.55s/it]
train_loss: 12.092806587092081
100%|██████████| 3750/3750 [1:36:05<00:00, 1.54s/it]
train_loss: 5.903924596150716

MNIST数据集
MNIST包含70,000张28x28像素的手写数字灰度图像(60,000张训练图像和10,000张测试图像),是计算机视觉和机器学习领域最常用的基准数据集之一。
train-images-idx3-ubyte.gz: training set images (9912422 bytes)
train-labels-idx1-ubyte.gz: training set labels (28881 bytes)
t10k-images-idx3-ubyte.gz: test set images (1648877 bytes)
t10k-labels-idx1-ubyte.gz: test set labels (4542 bytes)
数据集下载
pip3 install torchvision
国外网络环境,执行模块下载,需网络环境很稳定
import torchvision
# 正常下载数据集
train_data = torchvision.datasets.MNIST(
root='./data_download',
train=True,
download=True)
Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz
Failed to download (trying next):
HTTP Error 404: Not Found
Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz
Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz to ./data_download/MNIST/raw/train-images-idx3-ubyte.gz
Failed download. Trying https -> http instead. Downloading http://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz to ./data_download/MNIST/raw/train-images-idx3-ubyte.gz
100%|██████████| 9912422/9912422 [01:14<00:00, 132678.51it/s]
Extracting ./data_download/MNIST/raw/train-images-idx3-ubyte.gz to ./data_download/MNIST/raw
Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz
Failed to download (trying next):
HTTP Error 404: Not Found
Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-labels-idx1-ubyte.gz
Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-labels-idx1-ubyte.gz to ./data_download/MNIST/raw/train-labels-idx1-ubyte.gz
100%|██████████| 28881/28881 [00:00<00:00, 698760.33it/s]
Extracting ./data_download/MNIST/raw/train-labels-idx1-ubyte.gz to ./data_download/MNIST/raw
Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz
Failed to download (trying next):
HTTP Error 404: Not Found
Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-images-idx3-ubyte.gz
Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-images-idx3-ubyte.gz to ./data_download/MNIST/raw/t10k-images-idx3-ubyte.gz
100%|██████████| 1648877/1648877 [00:01<00:00, 1496713.58it/s]
Extracting ./data_download/MNIST/raw/t10k-images-idx3-ubyte.gz to ./data_download/MNIST/raw
Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz
Failed to download (trying next):
HTTP Error 404: Not Found
Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-labels-idx1-ubyte.gz
Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-labels-idx1-ubyte.gz to ./data_download/MNIST/raw/t10k-labels-idx1-ubyte.gz
100%|██████████| 4542/4542 [00:00<00:00, 1300377.39it/s]
Extracting ./data_download/MNIST/raw/t10k-labels-idx1-ubyte.gz to ./data_download/MNIST/raw

手动下载
使用浏览器搜索地址下载
https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz
https://ossci-datasets.s3.amazonaws.com/mnist/train-labels-idx1-ubyte.gz
https://ossci-datasets.s3.amazonaws.com/mnist/t10k-images-idx3-ubyte.gz
https://ossci-datasets.s3.amazonaws.com/mnist/t10k-labels-idx1-ubyte.gz
gz文件转npy文件
import numpy as np
import gzip
import os
def load_mnist_images(filename):
with gzip.open(filename, 'rb') as f:
data = np.frombuffer(f.read(), dtype=np.uint8, offset=16) # 跳过16字节头部
return data.reshape(-1, 28, 28)
def load_mnist_labels(filename):
with gzip.open(filename, 'rb') as f:
data = np.frombuffer(f.read(), dtype=np.uint8, offset=8) # 跳过8字节头部
return data
# 假设已下载并解压到当前目录
file_paths = {
'train_images': './data_download/mnist/raw/train-images-idx3-ubyte.gz',
'train_labels': './data_download/mnist/raw/train-labels-idx1-ubyte.gz',
'test_images': './data_download/mnist/raw/t10k-images-idx3-ubyte.gz',
'test_labels': './data_download/mnist/raw/t10k-labels-idx1-ubyte.gz'
}
# 加载并保存为 .npy
for name, path in file_paths.items():
if name.startswith('train'):
data = load_mnist_images(path) if 'image' in name else load_mnist_labels(path)
np.save(f'{name}.npy', data)
print(f"已保存 {name}.npy")
已保存 train_images.npy
已保存 train_labels.npy

载入npy数据
import numpy as np
x_train = np.load("./data_book/x_train.npy")
y_train_label = np.load("./data_book/y_train_label.npy")
print(x_train.shape)
print(y_train_label.shape)
(60000, 28, 28)
(60000,)
将每张28x28图片像素写入EXCEL中查看
from openpyxl import Workbook
from tqdm import tqdm # 进度条工具
wb = Workbook()
ws = wb.active
import numpy as np
x_train = np.load("./data_book/x_train.npy")
y_train_label = np.load("./data_book/y_train_label.npy")
num = x_train.shape[0]
height = x_train.shape[1]
weight = x_train.shape[2]
#num = 1
for n in tqdm(range(num)):
label_num = y_train_label[n]
for h in range(height):
for w in range(weight):
ws.cell(row=h+1,column=w+1).value = x_train[n][h][w]
#print(f'{h}_{w}:Finish')
wb.save(f'./xlsx_mnist/{n}图片序列号_mnist_label{label_num}.xlsx')
print("MNIST每个图片转换成XlSX文件完成")
if __name__ == '__main__':


682

被折叠的 条评论
为什么被折叠?



