下载数据集
import torch
import torchvision
import torchvision.transforms as transforms
# 定义数据预处理的转换
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# 训练集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,download=True, transform=transform)
# 加载
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,shuffle=True, num_workers=2)
# 测试集
testset = torchvision.datasets.CIFAR10(root='./data', train=False,download=True, transform=transform)
# 加载
testloader = torch.utils.data.DataLoader(testset, batch_size=4,shuffle=False, num_workers=2)
构建神经网络
import torch.nn as nn
import torch.nn.functional as F
class Net(nn.Module):
def __init__(self):
super(Net,self).__init__()
# 定义第一层卷积层
self.c1 = nn.Conv2d(3 ,6 ,5)
self.c2 = nn.Conv2d(6 ,15, 5)
# 定义池化
self.pool = nn.MaxPool2d(2,2)
# 定义全连接
self.fc1 = nn.Linear(15 * 5 * 5,120)
self.fc2 = nn.Linear(120 , 84)
slef.fc3 = nn.Linear(84 , 10)
def forward(self):
# 卷积层 relu激活函数去负值 池化
x = self.pool(F.relu(self.c1(x)))
x = self.pool(F.relu(self.c2(x)))
# 扁平化 转换成x形状 适配 全连接层的输入
x = x.view(-1, 15 * 5 * 5)
# 全连接层 - relu激活函数 去负值
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = F.slef.fc3(x)
return x
net = Net()
训练模型
# 损失函数 与 梯度清零函数
import torch.optim as optim
losser = nn.CrossEntropyLoss()
optimer = optim.SGD(net.parameters(),lr=0.1)
running_loss = 0
# 训练次数
for epoch in range(2):
# 遍历训练集
for data in trainloader:
# 提取图片与标签
images, labels = data
# 梯度清理
optimer.zero_grad()
# 计算训练结果张量
outputs = net(images)
# 计算损失值
loss = losser(outputs,labels)
# 反向传播
loss.backward()
# 参数更新
optimer.step()
# 打印训练数据
running_loss += loss.item()
if (i + 1) % 2000 == 0:
print('[{}, {}] loss: {}'.format(epoch + 1, i + 1, running_loss / 2000))
running_loss = 0
print("训练完毕")
保存模型
torch.save(net.state_dict(),'data.pth')
预测模型
correct = 0 # 正确率
total = 0 # 样本总量
# 禁止梯度清零 (测试/开发)
with torch.no_grad():
# 遍历测试集
for data in testloader:
images,labels = data
outputs = net(images)
_,predicted = torch.max(outputs,1) # 预测
total += lables.size(0) # 总量
# sum()方法用于计算张量中所有元素的总和
# item()将张量值转换成标量值
correct += (predicted == labels).sum().item()
print("正确率(%):", 100 * correct / total) # 打印正确率
print("样本总量",total)
每个类别的准确率
class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))
with torch.no_grad():
for data in testloader:
images, labels = data
outpus = net(images)
_,predicted = torch.max(outpus,1)
c = (predicted == labels).squeeze() # 挤压维度:其中删除了维度为1的维度,返回true,false
for i in range(4):
label = labels[i]
class_correct[label] += c[i].item()
class_total[label] += 1
for i in range(10):
print(classes[i], 100 * class_correct[i] / class_total[i])