环境为python3.8,如果系统支持gpu加速的话,可以把数据上传到gpu加速,我这里没用gpu,直接在cpu上训练。gpu加速详细操作不做赘述。
数据库用的经典的MINIST库,包含60000张训练图片和10000张测试图片。
上代码:(含注释)
# coding:utf-8
import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
#数据预处理
transform=transforms.Compose(
[transforms.ToTensor(),
transforms.Lambda(lambda x:x.repeat(3,1,1)),#黑白灰图片转化为rgb图片
transforms.Normalize(mean=(0.5,0.5,0.5),std=(0.5,0.5,0.5))]
)
#下载mnist库数据并定义读取器
train_set=torchvision.datasets.MNIST(root='./mnist',train=True,transform=transform,download=True)
train_loader=torch.utils.data.DataLoader(train_set,batch_size=100,shuffle=True,num_workers=0)
test_set=torchvision.datasets.MNIST(root='./mnist',train=False,transform=transform,download=True)
test_loader=torch.utils.data.DataLoader(test_set,batch_size=100,shuffle=True,num_workers=0)
'''
print(train_set.data.size())
print(train_set.targets.size())
print(test_set.data.size())
print(test_set.targets.size())
print(train_loader.dataset.data.shape)
print(test_loader.dataset.data.shape)
输出:
torch.Size([60000, 28, 28])
torch.Size([60000])
torch.Size([10000, 28, 28])
torch.Size([10000])
torch.Size([60000, 28, 28])
torch.Size([10000, 28, 28])
#打印图像尺寸与训练集、测试集数量
'''
#定义神经网络类
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1=nn.Sequential(
nn.Conv2d(3,10,5),#输入通道(rgb3个通道),输出通道,卷积核
nn.ReLU(),#激活函数
nn.MaxPool2d(2,2)#最大池化层输出
)
self.conv2=nn.Sequential(
nn.Conv2d(10,20,5),
nn.ReLU(),
nn.MaxPool2d(2,2)
)
self.fc1=nn.Sequential(
nn.Linear(20*4*4,10),
)
def forward(self,x):
x=self.conv1(x)
x1=self.conv2(x)
x1=x1.view(-1,20*4*4)
x2=self.fc1(x1)
return x2
net=Net()
#损失函数和优化器
criterion=nn.CrossEntropyLoss()
optimizer=torch.optim.Adam(net.parameters(),lr=0.001)
print("------------training-------------")
for epoch in range(2):#训练2次
running_loss = 0.0
for i,data in enumerate(train_loader,0):#每次训练100个
inputs,labels=data
optimizer.zero_grad()
outputs=net(inputs)
loss=criterion(outputs,labels)#计算损失函数
loss.backward()#自动反向传递
optimizer.step()#更新参数
running_loss += loss.item()
if i % 10 == 9:#每1000个打印一次损失函数
print('epoch:%d times:%5d progress[%d/%d] loss:%.3f' % (epoch + 1, i + 1,(i+1)*100,60000, running_loss / 1000))
running_loss = 0.0
print("Finished Training.")
import matplotlib.pyplot as plt
import numpy as np
def imshow(image):
image=image*0.5+0.5
npimg=image.numpy()#将Tensor变量转换为ndarray变量
plt.imshow(np.transpose(npimg,(1,2,0)))#显示灰度图片
plt.show()
print("-----------random test-----------")
print("一次随机抽取100个,打印前十个的答案和预测值:")
images,labels=next(iter(test_loader))
#print("images.size():",images.size(),"\nlabels:",labels)
imshow(torchvision.utils.make_grid(images))#make_grid的作用是将若干幅图像拼成一幅图像,这里显示四张合并为一张的图片,这里传入的还是tensor对象
for i in range(10):
print('GroundTruth:',labels[i].item())
outputs=net(images)
#print("outputs.size():",outputs.size())
_,predicted=torch.max(outputs,1)#dim=1表示输出所在行的最大值,即最大的特征值,_表示的是最大值的数值,predicted表示最大特征的标志
for i in range(10):
print('Predicted:',predicted[i].item())
print("-----------all test------------")
correct=0
total=0
with torch.no_grad():
for data in test_loader:
images,labels=data
outputs=net(images)
_,predicted=torch.max(outputs.data,1)
total+=labels.size(0)
correct+=(predicted==labels).sum().item()
print("总判断正确率: %d %%"%(100*correct/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 test_loader:
images,labels=data
outputs=net(images)
_, predicted = torch.max(outputs, 1)
c=(predicted==labels).squeeze()#根据预测结果和真实标签是否相等,输出 1 或者 0,数据类型为tensor
for i in range(100):
label=labels[i]
class_total[label]+=1
class_correct[label]+=c[i].item()
for i in range(10):
print("%d的预测正确率为: %.1f%%"%(i,100*class_correct[i]/class_total[i]))