PyTorch MNIST 实现
概述
MNIST 包含 0~9 的手写数字, 共有 60000 个训练集和 10000 个测试集. 数据的格式为单通道 28*28 的灰度图.
获取数据
def get_data():
"""获取数据"""
# 获取测试集
train = torchvision.datasets.MNIST(root="./data", train=True, download=True,
transform=torchvision.transforms.Compose([
torchvision.transforms.ToTensor(), # 转换成张量
torchvision.transforms.Normalize((0.1307,), (0.3081,)) # 标准化
]))
train_loader = DataLoader(train, batch_size=batch_size) # 分割测试集
# 获取测试集
test = torchvision.datasets.MNIST(root="./data", train=False, download=True,
transform=torchvision.transforms.Compose([
torchvision.transforms.ToTensor(), # 转换成张量
torchvision.transforms.Normalize((0.1307,), (0.3081,)) # 标准化
]))
test_loader = DataLoader(test, batch_size=batch_size) # 分割训练
# 返回分割好的训练集和测试集
return train_loader, test_loader
网络模型
class Model(torch.nn.Module):
def __init__(self):
super(Model, self).__init__()
# 卷积层
self.conv1 = torch.nn.Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1))
self.conv2 = torch.nn.Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1))
# Dropout层
self.dropout1 = torch.nn.Dropout(0.25)
self.dropout2 = torch.nn.Dropout(0.5)
# 全连接层
self.fc1 = torch.nn.Linear(9216, 128)
self.fc2 = torch.nn.Linear(128, 10)
def forward(self, x):
"""前向传播"""
# [b, 1, 28, 28] => [b, 32, 26, 26]
out = self.conv1(x)
out = F.relu(out)
# [b, 32, 26, 26] => [b, 64, 24, 24]
out = self.conv2(out)
out = F.relu(out)
# [b, 64, 24, 24] => [b, 64, 12, 12]
out = F.max_pool2d(out, 2)
out = self.dropout1(out)
# [b, 64, 12, 12] => [b, 64 * 12 * 12] => [b, 9216]
out = torch.flatten(out, 1)
# [b, 9216] => [b, 128]
out = self.fc1(out)
out = F.relu(out)
# [b, 128] => [b, 10]
out = self.dropout2(out)
out = self.fc2(out)
output = F.log_softmax(out, dim=1)
return output
train 函数
def train(model, epoch, train_loader):
"""训练"""
# 训练模式
model.train()
# 迭代
for step, (x, y) in enumerate(train_loader):
# 加速
if use_cuda:
model = model.cuda()
x, y = x.cuda(), y.cuda()
# 梯度清零
optimizer.zero_grad()
output = model(x)
# 计算损失
loss = F.nll_loss(output, y)
# 反向传播
loss.backward()
# 更新梯度
optimizer.step()
# 打印损失
if step % 50 == 0:
print('Epoch: {}, Step {}, Loss: {}'.format(epoch, step, loss))
test 函数
def test(model, test_loader):
"""测试"""
# 测试模式
model.eval()
# 存放正确个数
correct = 0
with torch.no_grad():
for x, y in test_loader:
# 加速
if use_cuda:
model = model.cuda()
x, y = x.cuda(), y.cuda()
# 获取结果
output = model(x)
# 预测结果
pred = output.argmax(dim=1, keepdim=True)
# 计算准确个数
correct += pred.eq(y.view_as(pred)).sum().item()
# 计算准确率
accuracy = correct / len(test_loader.dataset) * 100
# 输出准确
print("Test Accuracy: {}%".format(accuracy))
main 函数
def main():
# 获取数据
train_loader, test_loader = get_data()
# 迭代
for epoch in range(iteration_num):
print("\n================ epoch: {} ================".format(epoch))
train(network, epoch, train_loader)
test(network, test_loader)
完整代码
完整代码:
import torch
import torchvision
import torch.nn.functional as F
from torch.utils.data import DataLoader
class Model(torch.nn.Module):
def __init__(self):
super(Model, self).__init__()
# 卷积层
self.conv1 = torch.nn.Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1))
self.conv2 = torch.nn.Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1))
# Dropout层
self.dropout1 = torch.nn.Dropout(0.25)
self.dropout2 = torch.nn.Dropout(0.5)
# 全连接层
self.fc1 = torch.nn.Linear(9216, 128)
self.fc2 = torch.nn.Linear(128, 10)
def forward(self, x):
"""前向传播"""
# [b, 1, 28, 28] => [b, 32, 26, 26]
out = self.conv1(x)
out = F.relu(out)
# [b, 32, 26, 26] => [b, 64, 24, 24]
out = self.conv2(out)
out = F.relu(out)
# [b, 64, 24, 24] => [b, 64, 12, 12]
out = F.max_pool2d(out, 2)
out = self.dropout1(out)
# [b, 64, 12, 12] => [b, 64 * 12 * 12] => [b, 9216]
out = torch.flatten(out, 1)
# [b, 9216] => [b, 128]
out = self.fc1(out)
out = F.relu(out)
# [b, 128] => [b, 10]
out = self.dropout2(out)
out = self.fc2(out)
output = F.log_softmax(out, dim=1)
return output
# 定义超参数
batch_size = 64 # 一次训练的样本数目
learning_rate = 0.0001 # 学习率
iteration_num = 5 # 迭代次数
network = Model() # 实例化网络
print(network) # 调试输出网络结构
optimizer = torch.optim.Adam(network.parameters(), lr=learning_rate) # 优化器
# GPU 加速
use_cuda = torch.cuda.is_available()
print("是否使用 GPU 加速:", use_cuda)
def get_data():
"""获取数据"""
# 获取测试集
train = torchvision.datasets.MNIST(root="./data", train=True, download=True,
transform=torchvision.transforms.Compose([
torchvision.transforms.ToTensor(), # 转换成张量
torchvision.transforms.Normalize((0.1307,), (0.3081,)) # 标准化
]))
train_loader = DataLoader(train, batch_size=batch_size) # 分割测试集
# 获取测试集
test = torchvision.datasets.MNIST(root="./data", train=False, download=True,
transform=torchvision.transforms.Compose([
torchvision.transforms.ToTensor(), # 转换成张量
torchvision.transforms.Normalize((0.1307,), (0.3081,)) # 标准化
]))
test_loader = DataLoader(test, batch_size=batch_size) # 分割训练
# 返回分割好的训练集和测试集
return train_loader, test_loader
def train(model, epoch, train_loader):
"""训练"""
# 训练模式
![img](https://img-blog.csdnimg.cn/img_convert/bb90bab7a121c343a71719cb2c777ced.png)
![img](https://img-blog.csdnimg.cn/img_convert/228a537f658534b230b199ad837f5ae9.png)
**网上学习资料一大堆,但如果学到的知识不成体系,遇到问题时只是浅尝辄止,不再深入研究,那么很难做到真正的技术提升。**
**[需要这份系统化资料的朋友,可以戳这里获取](https://bbs.csdn.net/topics/618545628)**
**一个人可以走的很快,但一群人才能走的更远!不论你是正从事IT行业的老鸟或是对IT行业感兴趣的新人,都欢迎加入我们的的圈子(技术交流、学习资源、职场吐槽、大厂内推、面试辅导),让我们一起学习成长!**
[外链图片转存中...(img-LfpKXG0d-1714632763564)]
[外链图片转存中...(img-LmzKw9D2-1714632763564)]
**网上学习资料一大堆,但如果学到的知识不成体系,遇到问题时只是浅尝辄止,不再深入研究,那么很难做到真正的技术提升。**
**[需要这份系统化资料的朋友,可以戳这里获取](https://bbs.csdn.net/topics/618545628)**
**一个人可以走的很快,但一群人才能走的更远!不论你是正从事IT行业的老鸟或是对IT行业感兴趣的新人,都欢迎加入我们的的圈子(技术交流、学习资源、职场吐槽、大厂内推、面试辅导),让我们一起学习成长!**