1. 设计网络结构:
采用最原始的Lenrt5:
2. 代码编写:
import torch
from torch import nn
class Lenet5(nn.Module):
def __init__(self):
super(Lenet5, self).__init__()
# 卷积单元
self.conv_unit = nn.Sequential(
# x:[b, 3, 32, 32]
nn.Conv2d(3, 6, kernel_size=5, stride=1, padding=0),
nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
nn.Conv2d(6, 16, kernel_size=5, stride=1, padding=0),
nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
)
# 全连接单元
self.fc_unit = nn.Sequential(
nn.Linear(16*5*5, 120), # 全连接之前有一个打平操作,所以输入为三者之积
nn.ReLU(),
nn.Linear(120, 84),
nn.ReLU(),
nn.Linear(84, 10)
)
def forward(self, x):
# [b, 3, 32, 32] => [b, 16, 5, 5]
batchsz = x.size(0)
x = self.conv_unit(x)
# [b, 16, 5, 5] => [b, 16*5*5]
x = x.view(batchsz, -1)
# [b, 16*5*5] => [b, 10]
logits = self.fc_unit(x)
return logits
3.读取数据,进行训练,测试
import torch
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision import transforms
from lenet5 import Lenet5
from torch import nn, optim
def main():
batchsz = 32
# 读取数据,位置,是否为训练集,数据转化
cifar_train = datasets.CIFAR10("cifar", True, transform=transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor()]), download=True)
# 批次读取, 原数据,批次大小,是否乱序
cifar_train = DataLoader(cifar_train, batch_size=batchsz, shuffle=True)
cifar_test = datasets.CIFAR10("cifar", False, transform=transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor()]), download=True)
cifar_test = DataLoader(cifar_test, batch_size=batchsz, shuffle=True)
x, label = iter(cifar_train).next() # 迭代器,一条一条的获取数据
print("x, label", x.shape, label.shape)
device = torch.device("cuda") # 设置加速器
model = Lenet5().to(device) # 在载入模型
print(model)
criteon = nn.CrossEntropyLoss().to(device) # 计算交叉熵
optimizer = optim.Adam(model.parameters(), lr=1e-3) # 设置优化器
model.train()
for epoch in range(1000):
for batchid, (x, label) in enumerate(cifar_train):
# x:[b, 3, 32, 32]
# label:[b]
x, label = x.to(device), label.to(device) # 转化到cuda上
# logits:[b, 10]
# label:[b]
logits = model(x) # 计算输出
loss = criteon(logits, label) # 计算损失
# backprob
optimizer.zero_grad() # 梯度清零
loss.backward() # loss更新梯度
optimizer.step() # 前向传播,更新权重
print(epoch, loss.item())
model.eval()
with torch.no_grad():
total_crrect = 0
total_num = 0
for x, label in cifar_test:
x, label = x.to(device), label.to(device)
logits = model(x)
pred = logits.argmax(dim=1) # 返回最大值所在位置
total_crrect += torch.eq(pred, label).float().sum().item() # 判断与label是否相等,返回正确个数
total_num += x.size(0)
acc = total_crrect / total_num # 计算正确率
print(epoch, acc)
if __name__ == '__main__':
main()