# (1)加载必要的库
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
# (2)定义超参数
BATCH_SIZE = 16 # 每批处理的数据
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 用GPU还是CPU训练
EPOCHS = 10 # 训练数据集的轮次
# (3)构建pipeline,对图像做处理
pipeline = transforms.Compose([
transforms.ToTensor(), # /将图片转化成tensor
transforms.Normalize((0.1307,), (0.3081,)) # 正则化:降低模型的复杂化
])
# (4)下载、加载数据
from torch.utils.data import DataLoader
# 下载数据集
train_set = datasets.MNIST('data', train=True, download=True, transform=pipeline)
test_set = datasets.MNIST('data', train=False, download=True, transform=pipeline)
# 加载数据
train_loader = DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(test_set, batch_size=BATCH_SIZE, shuffle=True)
# (5) 构建网络模型
class Digit(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 10, 5) # 卷积层 1:灰度图片的通道, 10:输出通道, 5: 卷积核
self.conv2 = nn.Conv2d(10, 20, 3) # 卷积层 10:输入通道,20 :输出通道,3:卷积核大小
self.fc1 = nn.Linear(20 * 10 * 10, 500) # 全连接层 20*10*10:输入通道,500:输出通道
self.fc2 = nn.Linear(500, 10) # 全连接层 500:输入通道,10:输出通道
def forward(self, x): # 前向传播
input_size = x.size(0) # batch_size BATCH_SIZE = 16 # 每批处理的数据
x = self.conv1(x) # 第一个卷积层 输入:batch*1*28*28 , 输出:batch*10*24*24(28-5+1 = 24)
x = F.relu(x) # 激活函数 保持shape不变,输出:batch*10*24*24
x = F.max_pool2d(x, 2, 2) # 池化层 最大池化 输入:batch*10*24*24 输出:batch*10*12*12
x = self.conv2(x) # 第二个卷积层 输入:batch*10*12*12 输出:batch*20*10*10(12-3+1=10)
x = F.relu(x) # 激活函数
x = x.view(input_size, -1) # 拉平,把二维图片拉成一维度的,-1:自动计算维度:20*20*10 = 2000
x = self.fc1(x) # 全连接层 输入:batch*2000 输出:batch*500
x = F.relu(x) # 激活层,保持shpae不变
x = self.fc2(x) # 全连接层2 输入:batch*500 输出:batch*10
output = F.log_softmax(x, dim=1) # 计算分类就,每个数字的概率值
return output # 把概率值给返回回去
# (6) 定义优化器
model = Digit().to(DEVICE) # 创建一个模型,并把模型放在设备上面去
optimizer = optim.Adam(model.parameters()) # 定义一个优化器,更新模型的参数。优化器有很多种,这只是其中的一个
# (7)定义训练方法
def train_model(model, device, train_loader, optimizer,epoch):
# model:模型, device:运行模型的设备(cpu还是gpu), optimizer:优化器,epoch:训练集训练的轮数 EPOCHS = 10 # 训练数据集的轮次
# 模型训练
model.train()
for batch_index, (data, target) in enumerate(train_loader):
# 部署到DEVICE上去
data, target = data.to(device), target.to(device)
# 梯度初始化为0
optimizer.zero_grad()
# 训练后的结果
output = model(data)
# 计算损失
loss = F.cross_entropy(output, target) # 此处使用的是交叉熵损失,交叉熵损失适用于多分类的情况
# 找到概率值最大的下标
pred = output.max(1, keepdim=True) # pred = output.argmax(dim=1) 1:表示维度,这里是指横轴
# 反向传播
loss.backward()
# 参数优化 上一步的反向传播哦,不进行参数更新吗?
optimizer.step()
if batch_index % 3000 == 0: # 每处理3000张图片打印一次损失函数
print('Train Epoch : {} \t Loss : {:.6}'.format(epoch, loss.item())) # 输出的数字是6位
# (8)定义测试方法
def test_model(model, device, test_loader):
# 模型验证 验证什么?
model.eval()
# 正确率
correct = 0.0
# 测试损失 测试损失?这个是做什么用的?
test_loss = 0.0
with torch.no_grad(): # 不会计算梯度,也不会进行反向传播
for data, target in test_loader:
# 部署到device上
data, target = data.to(device), target.to(device)
# 测试数据
output = model(data)
# 计算测试损失
test_loss = F.cross_entropy(output, target).item()
# 找到概率值最大的下标
pred = output.max(1, keepdim=True)[1] # 值,索引
# 累计正确的值
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
print('Test -- Average loss : {:.4f}, Accuracy : {:.3f}\n'.format(test_loss,
100 * correct / len(test_loader.dataset)))
# (9)调用方法(7)(8)
for epoch in range(1, EPOCHS + 1):
train_model(model, DEVICE, train_loader, optimizer, epoch)
test_model(model, DEVICE, test_loader)