Pytorch之CIFAR10分类卷积神经网络( 二 )


(2)GPU训练
在使用GPU训练时,只需要将上面的代码作一些微调就可以了 。将cpu转换为gpu上,需要改三个地方:网络模型、损失函数、训练中加载的数据集和测试中加载的数据集
# 定义训练的设备device = torch.device("cuda")# 创建网络模型net = Cifarnet()net = net.to(device)# 训练集和测试集一样的操作imgs, targets = dataimgs = imgs.to(device)targets = targets.to(device)
四、测试与评估
# 开始测试net.eval()total_test_loss = 0total_accuracy = 0with torch.no_grad():for data in test_dataloader:imgs, targets = dataoutputs = net(imgs)loss = loss_fn(outputs, targets)total_test_loss += loss.item()# argmax找出每一张预测概率最大的下标accuracy = (outputs.argmax(1) == targets).sum()total_accuracy += accuracyprint("测试集的Loss:{}".format(total_test_loss))print("测试集的正确率:{}".format(total_accuracy/test_data_size))writer.add_scalar("test_loss", total_test_loss,total_test_step)writer.add_scalar("test_accurary", total_accuracy/test_data_size, total_test_step)total_test_step += 1
以下是使用CPU训练了5个epoch,所以数据并不是很理想,电脑配置高的建议100以上训练 。
测试集正确率

Pytorch之CIFAR10分类卷积神经网络

文章插图
测试集损失函数:
训练集损失函数:
使用 colab gpu进行训练200个epoch:
训练200次正确率也不是很高,可能是由于没有调参的问题 。
模型的保存,为后续实现分类作准备 。
# 保存模型torch.save(net, "cifdarnet.pth")print("模型已保存")
五、实现分类
利用刚刚训练完成的网络模型,进行图像的分类 。
既然是为了实现图像分类,那么我们就要先获得图片对应的标签(分类类别),在中可以在数据集预处理的时候通过添加断点Debug的模式获取 。
# ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']class_id = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
注意:如果使用GPU训练的模型,而使用CPU去加载测试,需要做如下处理:
将gpu改为cpu时,遇到一个报错:
:toon a CUDAbut torch.cuda.() is False. If you areon a CPU-only ,use torch.load with ='cpu' to map yourto the CPU.
model = torch.load("cifdarnet_gpu.pth", map_location='cpu')
(1)以图片的形式测试
获取到类别后,然后就是将读取的图片进行预处理,其大小和信息要符合我们前面定义的网络模型 。
# image_path = 'D:\Data\Learn_Pytorch\cifar10_train\dog.png'image_path = 'D:\Data\Learn_Pytorch\cifar10_train\plane.png'image_raw = Image.open(image_path)image = image_raw.convert('RGB')# 将ARGB-->RGBprint(image)# [3,32,32]# 网络的输入为32X32transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32, 32)),torchvision.transforms.ToTensor()])image = transform(image)print(image.shape)
预处理主要的操作有颜色空间类型转换,要符合RGB颜色空间,其次图片的缩放,图片的尺寸要符合32x32,并且图片的类型要是类型 。
然后就可以加载前面保存的网络模型进行图像的分类:
model = torch.load("cifdarnet.pth")print(model)# [3,32,32] --> [1,3,32,32]image = torch.reshape(image,[1, 3, 32, 32])model.eval()with torch.no_grad():output = model(image)idx = output.argmax(1)reslut = class_id[idx]# 可视化draw = ImageDraw.Draw(image_raw)draw.text((10, 10), reslut, fill=(255, 0, 0))image_raw.show()print(output)