【PyTorch官方教程中文版学习笔记05】PyTorch 图像分类器--细节完善&尝试调用本地图片测试
- 网络模型的读取和存储
通过存储，存储训练好的神经网络参数。
通过加载读取，直接使用训练好的神经网络参数。
#保存方式1，以VGG为例（官方推荐的保存方式） torch.save(vgg16.state_dict(),"vgg16_method2.pth")#vgg16_method2.pth为自定义，后缀.pth #加载模型 model = torch.load("vgg16_method2.pth") print(model) #结果为字典形式的模型参数 vgg16 = torchvision.models.vgg16(pretrained=false)#网络模型结构 vgg16.load_state_dict( torch.load("vgg16_method2.pth"))#恢复成网络模型
#保存方式2，以Net为例 torch.save(net,"net_{}.pth".format(epoch)) #加载方式 class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(3, 6, 5) self.pool = nn.MaxPool2d(2, 2) self.conv2 = nn.Conv2d(6, 16, 5) self.fc1 = nn.Linear(16 * 5 * 5, 120) self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, 10) def forward(self, x): x = self.pool(F.relu(self.conv1(x))) x = self.pool(F.relu(self.conv2(x))) x = x.view(-1, 16 * 5 * 5) x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x #net 的类 model = torch.load("net_0.pth")
2. 神经网络转到GPU上跑（前提是有GPU）
需要把：
网络模型转移到cuda-->net = net.cuda() 损失函数转移到cuda-->criterion = criterion.cuda() 数据转移到cuda--> inputs = inputs.cuda() labels = labels.cuda()#或者可以用to(device)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") net.to(device)
criterion =.to(device)
inputs, labels = inputs.to(device), labels.to(device)
添加位置如下：
import torch import torch.nn as nn import torch.nn.functional as F import torchvision import torchvision.transforms as transforms import torch.optim as optim import time#用于计算时间 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=False, transform=transform) trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=False, num_workers=2) #测试集 testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=False, transform=transform) testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=2) #标签数组 classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') #随机收取四个样本 dataiter = iter(trainloader) images, labels = dataiter.next()
image = image.cuda()
labels =image.cuda()#这里也要移到gpu上 print(' '.join('%5s' % classes[labels[j]] for j in range(4))) #搭建网络Net class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(3, 6, 5) self.pool = nn.MaxPool2d(2, 2) self.conv2 = nn.Conv2d(6, 16, 5) self.fc1 = nn.Linear(16 * 5 * 5, 120) self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, 10) def forward(self, x): x = self.pool(F.relu(self.conv1(x))) x = self.pool(F.relu(self.conv2(x))) x = x.view(-1, 16 * 5 * 5) x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x net = Net() net = net.cuda() #网络模型转移到cuda #定义损失函数 criterion = nn.CrossEntropyLoss() criterion = criterion.cuda() #损失函数转移到cuda optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9) #开始10轮训练 start_time = time.time()#开始时间 for epoch in range(2): running_loss = 0.0 for i, data in enumerate(trainloader, 0): inputs, labels = data inputs = inputs.cuda() labels = labels.cuda() #数据转移到cuda outputs = net(inputs) loss = criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() running_loss += loss.item() if i % 100 == 0: print('[%d, %5d] loss: %.3f' %(epoch + 1, i , running_loss / 100)) end_time = time.time()#结束时间 print("100 mini-batches训练时间 : " + str(end_time-start_time))#输出用时 running_loss = 0.0 correct = 0 total = 0 with torch.no_grad(): for data in testloader: inputs, labels = data inputs = inputs.cuda() labels = labels.cuda() #数据转移到cuda outputs = net(inputs) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() print('Accuracy of the network on the 10000 test images: %d %%' % (100 * correct / total)) torch.save(net,"net_{}.pth".format(epoch)) print("模型已保存") print('Finished Training') #10轮训练后的Net模型对图像识别的精确度 outputs = net(images) _, predicted = torch.max(outputs.data, 1) print('Predicted: ', ' '.join('%5s' % classes[predicted[j]] for j in range(4)))
GPU和CPU上运行速度的对比：
为什么没有注意到与CPU相比巨大的加速？因为你的网络非常小。 3. 调用本地文件--相对路径image_path = "img/cat.jpg"#同一层级文件,其中cat.jpg是在网上随机找的图片 image = Image.open(image_path) print(image)
#输出原图片的属性，可见size不符合神经网络的shape输入要求#根据神经网络的输入的shape，对图像进行transform from PIL import Image image = image.convert("RGB")#jpg/png的图片是四通道，除了三个颜色通道外，还有一个透明通道
transform = transforms.Compose( [transforms.Resize((32,32)),transforms.ToTensor()])
#size变成32*32，转换一个PIL库的图片或者numpy的数组为tensor张量类型 image = transform(image)
#如果载入的神经网络是在GPU上完成训练的，那么调用的本地图像也要转到GPU上才能进行测试 image = image.cuda()
完整程序：
整个框架使用了上一章的神经网络，并调用了第五次的训练参数。因为5次的训练还远远不够，所以精准度不是很高，只能识别特征特别明显的图片。
import torch import torch.nn as nn import torch.nn.functional as F import torchvision.transforms as transforms from PIL import Image classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') image_path = "img/cat.jpg"#输入一张猫猫的靓照 image = Image.open(image_path) image = image.convert("RGB")#jpg/png的图片是四通道，除了三个颜色通道外，还有一个透明通道 transform = transforms.Compose( [transforms.Resize((32,32)),transforms.ToTensor()]) image = transform(image) image = image.cuda() #拷贝网络模型 class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(3, 6, 5) self.pool = nn.MaxPool2d(2, 2) self.conv2 = nn.Conv2d(6, 16, 5) self.fc1 = nn.Linear(16 * 5 * 5, 120) self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, 10) def forward(self, x): x = self.pool(F.relu(self.conv1(x))) x = self.pool(F.relu(self.conv2(x))) x = x.view(-1, 16 * 5 * 5) x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x model = torch.load("net_4.pth") output = model(image) print("输出结果：") print(classes[output.argmax(1)])
输出：
结果一致！十分开心，入门课程完成撒花花?！！