Pytorch快速上手

一、加载数据集（Dataset）

加载数据集需要继承Dataset，通常情况下需要实现__init__方法、__getitem__方法以及__len__方法。

案例一：

import osimport torch
from torch.utils.data import Dataset
from PIL import Imageclass MyData(Dataset):'''读取数据集中的数据'''def __init__(self, root_dir, lable_dir):'''初始化加载数据，得到所有图片的名称:param root_dir: 目录路径:param lable_dir: 标签名'''self.root_dir = root_dirself.lable_dir = lable_dirself.path = os.path.join(self.root_dir, self.lable_dir)   # 拼接路径得到具体的图片路径self.img_path = os.listdir(self.path)   # 对应数据集中的所有文件名称def __getitem__(self, item):'''重写迭代的方式，这样做的好处是可以在加载数据的时候能够直接遍历到每一个图片，以及对应的标签:param item: 迭代的索引:return: 返回的图片以及图片对应的标签'''img_name = self.img_path[item]img_item_path = os.path.join(self.root_dir, self.lable_dir, img_name)   # 得到文件的路径img = Image.open(img_item_path)  # 读入图片文件label = self.lable_dir  # 得到图片类型return img, labeldef __len__(self):''':return: 返回图片文件的个数'''return len(self.img_path)if __name__ == '__main__':root_dir = r"..\datasets\train"cats_label_dir = "cats"dogs_label_dir = "dogs"cats_dataset = MyData(root_dir, cats_label_dir)  # 得到猫的图片集合# for img, label in cats_dataset:#     print(label)  # 输出标签#     img.show()  # 展示所有图片dogs_dataset = MyData(root_dir, dogs_label_dir)  # 得到狗的图片集合train_dataset = cats_dataset + dogs_dataset   # 将数据集进行拼接，得到整个训练集的图片集合

二、显示训练的情况（`tensorboard`）

tensorboard是一个可视化的包，可用于展示模型损失的变化或者是直接查看图片数据集。

要使用tensorboard，需要额外进行下载，并且版本过高可能会出现异常。

pip install tensorboard==2.12.0

查看结果需要在当前项目的控制台上输入：

tensorboard --logdir=logs --port=6007  
参数--logdir的值是对应的生成的文件的路径。
参数--port的值表示打开的端口号，可以自定义（避免与其它端口号产生冲突）

运行之后会出现如下的信息，直接点击蓝色显示的地址即可打开对应的网页查看具体的信息。

在这里插入图片描述

案例一：

from torch.utils.tensorboard import SummaryWriter
'''tensorboard需要进行额外下载：pip install tensorboard==2.12.0  版本过高可能会出现异常'''
writer = SummaryWriter(r'..\logs')  # 将生成的文件放入到logs文件夹中（若没有logs文件夹则自动创建）# writer.add_image()  # 添加图片
# writer.add_scalar()  # 添加标量，tag：标签， scalar_value：y轴， global_step：x轴
# writer.close()  # 关闭通道for i in range(100):writer.add_scalar('y = x', i, i)writer.close()'''
查看信息：
在当前控制台输入：tensorboard --logdir=logs --port=6007
'''

打开给出的地址（如：http://localhost:6007/）之后即可在相应的网页上查看具体的信息了。

案例一结果：

在这里插入图片描述

add_image的使用
    def add_image(tag, img_tensor, global_step=None, walltime=None, dataformats="CHW")
参数解释：

tag：数据表示符
img_tensor：图片数据（支持的数据类型：torch.Tensor, numpy.array, string, blobname）
global_step：全局步长值（表示同一个表示符下的数据）
walltime：Optional override default walltime (time.time())
dataformats：表示img_tensor中的每一维表示什么意思。

注：opencv读取的数据类型是numpy类型。

案例二：

from torch.utils.tensorboard import SummaryWriter
from PIL import Image
import numpy as npwriter = SummaryWriter('../logs')
image_path = '../datasets/train/cats/1.jpeg'
img = Image.open(image_path)
img_array = np.array(img)  # 将图片信息转化为numpy类型writer.add_image(tag="cat", img_tensor=img_array, global_step=1, dataformats="HWC")  # pytorch数据类型默认为[batch_size, C, W, H]for i in range(100):writer.add_scalar("y = 2x", 2 * i, i)
writer.close()

案例二结果：

在这里插入图片描述

二、数据变换（`Transforms`）

Transforms用于对数据进行变换，常见的有转换为tensor数据类型、对数据进行归一化、进行缩放，进行随机裁剪等。

# 导入transforms
from torchvision import transforms

由于transforms是在torchvision库中，并且这个库也需要进行额外下载。

# pip安装方式：
pip install torchvision
# conda安装方式：
conda install torchvision -c pytorch

转换为tensor数据类型：

tensor_trains = transforms.ToTensor()   # 首先实例化一个转换器
tensor_img = tensor_trains(img)   # 传入图片进行数据类型的转换

进行归一化：

trans_norm = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])  # 参数表示各个维度的均值和方差
img_norm = trans_norm(img)  # img需要是tensor类型的

进行缩放：

trans_resi = transforms.Resize((512, 512))  # 参数表示需缩放后的尺寸
img_resi = trans_resi(img)  # 对img进行等比缩放，img需要是tensor类型的

进行随机裁剪：

trans_randcop = transforms.RandomCrop(128)  # 表示裁剪的区域为128 * 128
img_randcop = trans_randcop(img)  # 对img进行随机裁剪，img需要是tensor类型的

组合：

trans_resi = transforms.Resize(512)  # 实例化一个缩放，参数为单个值表示与最短的边进行匹配
trans_norm = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])  # 实例化一个归一化器
trans_comp = transforms.Compose([trans_resi, trans_norm])  # 将两个实例放入到Compose中。
trans_comp(img)  # img使用Compose实例。

案例一:

转换为tensor数据类型

from PIL import Image
from torchvision import transforms# 构建tensor数据
img_path = "../datasets/train/cats/3.jpeg"
img = Image.open(img_path)
tensor_trains = transforms.ToTensor()  # 转换为tensor类型
tensor_img = tensor_trains(img)
# print(tensor_img)
print(type(tensor_img))

案例二：

转换为tensor类型，并使用tensorboard进行加载

from PIL import Image
from torch.utils.tensorboard import SummaryWriter
from torchvision import transformsimg_path = "../datasets/train/dogs/19.jpeg"
img = Image.open(img_path)
writer = SummaryWriter("../logs")tensor_train = transforms.ToTensor()
tensor_img = tensor_train(img)
writer.add_image("Tensor_img", tensor_img)
writer.close()

案例二结果：

在这里插入图片描述

案例三：

常用的transforms

'''常用的transform'''
from PIL import Image
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms
import numpy as np
writer = SummaryWriter('../logs')# 归一化
def Norm(img):trans_norm = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])  # 参数表示各个维度的均值和方差img_norm = trans_norm(img)  # img需要是tensor类型的return img_norm# 等比缩放
def Resi(img):trans_resi = transforms.Resize((512, 512))img_resi = trans_resi(img)return img_resi# 等比缩放且进行归一化
def Comp(img):'''参数不能是tensor类型的'''trans_resi = transforms.Resize(512)  # 单个值表示与最短的边进行匹配trans_norm = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])# trans_tensor = transforms.ToTensor()trans_comp = transforms.Compose([trans_resi, trans_norm])return trans_comp(img)# 进行随机裁剪
def RandCop(img):ls = []trans_randcop = transforms.RandomCrop(128)  # 表示裁剪的区域为128 * 128for i in range(10):img_randcop = trans_randcop(img)ls.append(img_randcop)return lsif __name__ == '__main__':img_path = '../datasets/train/cats/83.jpeg'img = Image.open(img_path)img_transform = transforms.ToTensor()img_tensor = img_transform(img)# 进行归一化writer.add_image("normalize", Norm(img_tensor))# 等比缩放writer.add_image("resize", Resi(img_tensor))# 等比缩放+归一化writer.add_image("norm+resize", Comp(img_tensor))# 随机裁剪for idx, imgt in enumerate(RandCop(img_tensor)):writer.add_image("randomCrop", imgt, idx)

案例三结果：

在这里插入图片描述

三、`torchvision`中数据集的使用

Pytorch中提供了一些标准的数据集，可在官方网页进行查看。

案例一：

import torchvision
from torch.utils.tensorboard import SummaryWriterdataset_transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])
train_set = torchvision.datasets.CIFAR10(root='../CIFAR10', train=True, transform=dataset_transform, download=True)
test_set = torchvision.datasets.CIFAR10(root='../CIFAR10', train=False, transform=dataset_transform, download=True)writer = SummaryWriter('../cir_log')
for i in range(10):img, target = test_set[i]writer.add_image("CIFAR10", img, i)
writer.close()

参数解释：
train_set = torchvision.datasets.CIFAR10(root='../CIFAR10', train=True, transform=dataset_transform, download=True)
root：表示数据集的位置（如果下载的话表示存放数据集的位置）
train：表示是否为训练数据，如果为True表示得到的是训练集，如果为False表示得到的是验证集
transform：表示对数据进行transform数据变换

四、`Dataloader`的使用

案例一：

import torchvision
# 准备测试数据集
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision.utils import make_gridtest_data = torchvision.datasets.CIFAR10("../CIFAR10", train=False, transform=torchvision.transforms.ToTensor())
test_loader = DataLoader(dataset=test_data, batch_size=64, shuffle=True, num_workers=0, drop_last=False)img, target = test_data[0]
print(img.shape)
print(target)
writer = SummaryWriter('../cir_log')
for epoc in range(2):step = 0for data in test_loader:imgs, targets = data# grid_img = make_grid(imgs, nrow=8)writer.add_image("EPOC{}".format(epoc), imgs, step, dataformats='NCHW')  # 指定batch_size的维度信息step += 1
writer.close()

参数解释：

dataset：表示输入的数据
batch_size：表示对多少个数据进行打包（一次性处理多少个数据）
shuffle：表示每轮读取样本时，是否进行随机打乱。
num_workers：使用多线程进行加载数据，默认值为0，表示只使用主线程进行读取数据。
drop_last：每一次读取数据时，如果最后的一些数据不能组成完整的一组是否进行抛弃，False表示不进行抛弃。

`nn.module`的使用

nn.module用于自定义自己的模型，可以方便快速的搭建自己的模型。

案例一

import torch
from torch import nnclass MyModule(nn.Module):def __init__(self):super(MyModule, self).__init__()   # 继承父类def forward(self, input):   # 编写forward方法，表示前向传播的过程output = input * inputreturn outputif __name__ == '__main__':module = MyModule()x = torch.tensor(2.0)out = module(x)print(out)

五、卷积

torch.nn.functional.F.conv2d(in_channels: int,out_channels: int,kernel_size: _size_2_t,stride: _size_2_t = 1,padding: Union[str, _size_2_t] = 0,dilation: _size_2_t = 1,groups: int = 1,bias: bool = True,padding_mode: str = 'zeros',  # TODO: refine this typedevice=None,dtype=None)

参数讲解：

in_channels：输入的特征图通道大小
kernel_size：卷积核大小
stride：卷积核移动的步长
padding：外围填充多少个0
dilation：用来设置卷积核的间隔（空洞卷积）
groups：卷积层的groups参数指定了将输入数据分为多少个组，并且每个组的通道数相同。这个参数主要用于实现分组卷积，通过将输入数据分为多个组来减少参数数量和计算量，从而提高网络的效率。在使用groups参数时，需要确保输入数据的通道数能够被groups参数整除。
bias：是否添加偏置项

案例一

'''卷积操作'''
import torch
import torch.nn.functional as Finput_tensor = torch.tensor([[1, 2, 0, 3, 1],[0, 1, 2, 3, 1],[1, 2, 1, 0, 0],[5, 2, 3, 1, 1],[2, 1, 0, 1, 1]])
kernel_tensor = torch.tensor([[1, 2, 1],[0, 1, 0],[2, 1, 0]])input_tensor = torch.reshape(input_tensor, (1, 1, 5, 5))   # 编写为标准的维度信息：[batch_size, channel, width, heigh]
kernel_tensor = torch.reshape(kernel_tensor, (1, 1, 3, 3))
print(input_tensor.shape)
print(kernel_tensor.shape)output1 = F.conv2d(input=input_tensor, weight=kernel_tensor, stride=1)
print("output1", output1)output2 = F.conv2d(input=input_tensor, weight=kernel_tensor, stride=2)
print("output2", output2)output3 = F.conv2d(input=input_tensor, weight=kernel_tensor, stride=1, padding=1)
print("output3", output3)

案例二

import torch
import torchvision
from torch import nn
from torch.nn import Conv2d
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriterdataset = torchvision.datasets.CIFAR10("../CIFAR10", train=False, transform=torchvision.transforms.ToTensor(),download=False)
dataloader = DataLoader(dataset, batch_size=64)class MyModule(nn.Module):def __init__(self):super(MyModule, self).__init__()self.conv1 = Conv2d(in_channels=3, out_channels=6, kernel_size=(3, 3), stride=(1, 1), padding=0)  # 定义一个卷层def forward(self, x):x = self.conv1(x)return xif __name__ == '__main__':module = MyModule()writer = SummaryWriter('../logs')for step, data in enumerate(dataloader):imgs, targets = dataoutput = module(imgs)print(imgs.shape)print(output.shape)writer.add_image("input", imgs, step, dataformats="NCHW")output = torch.reshape(output, (-1, 3, 30, 30))  # 降低通道数writer.add_image("output", output, step, dataformats="NCHW")

计算卷积后图像宽和高的公式

$Input：(N， C_{in}，H_{in}，W_{in})$
$Output：(N，C_{out}，H_{out}，W_{out})$

$H_{out} = [\frac{H_{in} + 2 \times padding[0] - dilation[0] \times (kernel_size[0] - 1) - 1}{stride[0]}]$

$W_{out} = [\frac{W_{in} + 2 \times padding[1] - dilation[1] \times (kernel_size[1] - 1) - 1}{stride[1]}]$

六、池化层

self.maxpool1 = MaxPool2d(kernel_size=(3, 3), ceil_mode=False)  # ceil_mode: 向下取整

参数解释：

kernel_size：卷积核大小
ceil_mode：是否进行向下取整，（向下取整表示：在卷积核移动过程中所扫过的图像如果不完整【如：（3 x 2）但卷积核是3 x 3】，此时是否保留最大值）

案例一

'''池化层'''
import torch
import torchvision
from torch import nn
from torch.nn import MaxPool2d
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriterdataset = torchvision.datasets.CIFAR10("../CIFAR10", train=False, transform=torchvision.transforms.ToTensor(),download=False)
dataloader = DataLoader(dataset, batch_size=64)class MyModule(nn.Module):def __init__(self):super(MyModule, self).__init__()self.maxpool1 = MaxPool2d(kernel_size=(3, 3), ceil_mode=False)  # ceil_mode: 向下取整def forward(self, x):output = self.maxpool1(x)return outputif __name__ == '__main__':module = MyModule()writer = SummaryWriter('../logs')for step, data in enumerate(dataloader):imgs, targets = datawriter.add_image("input_pool", imgs, step, dataformats="NCHW")output = module(imgs)writer.add_image("output_pool", output, step, dataformats="NCHW")writer.close()

七、非线性激活层

self.relu1 = ReLU()
self.sigmoid1 = Sigmoid()

案例一

'''非线性激活层'''
import torch
import torchvision
from torch import nn
from torch.nn import ReLU, Sigmoid
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriterdataset = torchvision.datasets.CIFAR10("../CIFAR10", train=False, transform=torchvision.transforms.ToTensor(),download=False)
dataloader = DataLoader(dataset, batch_size=64)class MyModule(nn.Module):def __init__(self):super(MyModule, self).__init__()self.relu1 = ReLU()self.sigmoid1 = Sigmoid()def forward(self, x):output = self.sigmoid1(x)return outputif __name__ == '__main__':writer = SummaryWriter('../logs')module = MyModule()for step, data in enumerate(dataloader):imgs, targets = datawriter.add_image("input_sigmod", imgs, step, dataformats="NCHW")output = module(imgs)writer.add_image("output_sigmod", output, step, dataformats="NCHW")writer.close()

八、线性激活层

线性层

torch.nn.Linear(in_features, out_features, bias=True,)

归一化层

torch.nn.BatchNorm2d(num_features, eps=1e-05, momentum-0.1, affine=True, track_running_status=True)

案例一

'''线性层以及其它层'''
import torch
import numpy as np
import torchvision
from torch import nn
from torch.nn import Linear
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriterdataset = torchvision.datasets.CIFAR10("../CIFAR10", train=False, transform=torchvision.transforms.ToTensor(),download=False)
dataloader = DataLoader(dataset, batch_size=64)class MyModule(nn.Module):def __init__(self):super(MyModule, self).__init__()self.linear1 = Linear(in_features=196608, out_features=10)  # in_features的值由ft = torch.flatten(imgs)计算得出def forward(self, x):output = self.linear1(x)return outputif __name__ == '__main__':module = MyModule()for step, data in enumerate(dataloader):imgs, targets = data# print(imgs.shape)# ft = torch.flatten(imgs)  # 展平处理ft = torch.reshape(imgs, (1, 1, 1, -1))if ft.size(3) != 196608:continue# print(ft.shape)  # 得到线性层输入的大小output = module(ft)print(output.shape)

九、损失(Loss)

L1损失(L1loss) ：平均绝对值误差（MAE）

将目标值与预测值作差求和再取平均值。

$L_n = \sum\limits_{i = 1}^n |x_i - y_i| / n$

torch.nn.L1Loss(size_average=None, reduce=None, reduction='mean')  # reduction也可以设置为sum

L2损失(L2loss)：欧氏距离(`MSELoss`)

$L_n = \frac{1}{n} \sum\limits_{i = 1}^n(x_i - y_i)^2$

torch.nn.MSELoss(size_average=None,reduce=None, reduction='mean')  # reduction也可以设置为sun

二元交叉熵损失(`Binary Crossentropy Loss`)

$L_n = -\sum\limits y_i \cdot \log{{p_i}}$ $y_i$ 表示真实标签， $p_i$ 表示模型预测出的概率值。

# 创建 CrossEntropyLoss 对象
criterion = nn.CrossEntropyLoss()
# 计算二元交叉熵损失
loss = criterion(pred, target)

交叉熵(CROSSENTROPYLOSS)

$-\log {\frac{exp(x[class])}{\sum\limits_j exp(x[j])}} = -x[class] + \log{\sum\limits_j exp(x[j])}$

torch.nn.CrossEntropyLoss(weight=None, size_average=None, ignore_index=-100, reduce=None, reduction='mean')

十、优化器

优化器用于更新参数，实际上是利用backward得到的梯度进行更新参数。

optimizer = optiom.SGD(model.paramenters(), lr=0.01, momentum=0.9)
optimizer = option.Adam([var1, var2], lr=0.001)   # 常用

十一、模型迁移

模型迁移是在已经训练好的模型上进行训练，也可以成为模型的微调，具体的做法有再原有的模型上添加层，或者修改原来模型的某些层。

import torchvision
from torch import nnvgg16_false = torchvision.models.vgg16()
vgg16_true = torchvision.models.vgg16(weights=torchvision.models.VGG16_Weights)  # 加载带权重的模型train_data = torchvision.datasets.CIFAR10("../CIFAR10", train=True, transform=torchvision.transforms.ToTensor(), download=False)
vgg16_true.add_module("add_linear",  nn.Linear(100, 10))  # 在模型最后添加一个线性层，第一个参数为线性层的名称，第二个参数为层的类型# 在容器中添加一层
vgg16_true.classifier.add_module("add_linear",  nn.Linear(1000, 10))# 修改某一个层
vgg16_true.classifier[6] = nn.Module(4096, 10)

十二、权重的加载与保存

import torch
import torchvision
vgg16 = torchvision.models.vgg16(weights=torchvision.models.VGG16_Weights)

第一种权重的加载与保存

# 方式一：（保存了模型的结构与参数）
# 1. 保存
torch.save(vgg16, "vgg16_method.pth")
# 2. 读取
module = torch.load("vgg16_method.pth")

第二种权重的加载与保存

# 方式二：（只保存了模型的参数，其格式为字典格式）
# 1. 保存
torch.save(vgg16.state_dict(), "vgg16_method.pth")
# 2. 读取
vgg16 = torchvision.models.vgg16()  # 先创建模型，在导入权重
vgg16.load_state_dict(torch.load("vgg16_method.pth"))

十三、使用GPU训练模型

第一种方式

找到网络模型、数据（输入、标注）、损失函数后加.cuda()即可。

具体的可以看下面的实例部分

健壮性更高的方法：

if torch.cuda.is_available():module = MyModule().cuda()

第二种方式

先指定设备

Device = torch.device("cpu")   # 指定设备为cpu
Device = torch.device("cuda")   # 指定设备为GPU
Device = torch.device("cuda:0")   # 指定设备为第一块GPU，若只有一块则与cuda等效
Device = torch.device("cuda:1")   # 指定设备为第二快GPU

再使用.to(Device)指定训练设备

十四、实例

model.py

模型程序

import torch
from torch import nnclass MyModule(nn.Module):def __init__(self):super(MyModule, self).__init__()self.model = nn.Sequential(nn.Conv2d(in_channels=3, out_channels=32, kernel_size=(5, 5), stride=(1, 1), padding=2),nn.MaxPool2d(kernel_size=(2, 2)),nn.Conv2d(in_channels=32, out_channels=32, kernel_size=(5, 5), stride=(1, 1), padding=2),nn.MaxPool2d(kernel_size=(2, 2)),nn.Conv2d(in_channels=32, out_channels=64, kernel_size=(5, 5), stride=(1, 1), padding=2),nn.MaxPool2d(kernel_size=(2, 2)),nn.Flatten(),nn.Linear(in_features=(64 * 4 * 4), out_features=64),nn.Linear(in_features=64, out_features=10))def forward(self, x):x = self.model(x)return x

train.py

训练程序

import torch
import torchvision
from torch.utils.tensorboard import SummaryWriter
from model import *
from torch import nn
from torch.utils.data import DataLoader# 得到数据集
train_data = torchvision.datasets.CIFAR10(root="../CIFAR10", train=True, transform=torchvision.transforms.ToTensor(), download=False)
test_data = torchvision.datasets.CIFAR10(root="../CIFAR10", train=False, transform=torchvision.transforms.ToTensor(), download=False)# 得到数据的数量
train_data_size = len(train_data)
test_data_size = len(test_data)print("训练数据集的长度为：{}".format(train_data_size))
print("测试数据集的长度为：{}".format(test_data_size))# 利用Dataloader来加载数据
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)# 创建网路
module = MyModule()
module = module.cuda()  # 加载到GPU上
# 损失函数
loss_fn = nn.CrossEntropyLoss()
# 优化器
learning_rate = 1e-2
optimizer = torch.optim.SGD(module.parameters(), lr=learning_rate)# 设置训练网路的一些参数
total_train_step = 0
total_test_step = 0
epoch = 10
writer = SummaryWriter('logs')for i in range(epoch):print("------------第{}轮训练开始-----------".format(i + 1))# 训练步骤开始module.train()for data in train_dataloader:imgs, targets = dataimgs = imgs.cuda()  # 加载到GPU上targets = targets.cuda()  # 加载到GPU上outputs = module(imgs)loss = loss_fn(outputs, targets)loss_fn = loss_fn.cuda()    # 加载到GPU上# 优化器进行优化optimizer.zero_grad()loss.backward()optimizer.step()total_train_step += 1if total_test_step % 100 == 0:print("训练次数{}， loss:{}".format(total_train_step, loss.item()))writer.add_scalar("train_loss", loss.item(), total_train_step)# 测试步骤开始module.eval()total_test_loss = 0total_accuracy = 0with torch.no_grad():  # 将梯度清零for data in test_dataloader:imgs, targets = dataimgs = imgs.cuda()   # 加载到GPU上targets = targets.cuda()  # 加载到GPU上outputs = module(imgs)loss = loss_fn(outputs, targets)total_test_loss = total_test_loss + loss.item()accuracy = (outputs.argmax(1) == targets).sum()total_accuracy = total_accuracy + accuracyprint("整体测试集上的loss：{}".format(total_test_loss))print("整体训练集上的loss：{}".format(total_accuracy / test_data_size))writer.add_scalar("test_loss", total_test_loss, total_test_step)writer.add_scalar("test_accuracy", total_accuracy / test_data_size, total_train_step)total_test_step = total_test_step + 1torch.save(module, "./weights/module_{}.pth".format(i))print("模型已保存")
writer.close()

verify.py

推理程序（验证程序）

import torchvisionfrom model import *
import torch
from PIL import Imageif __name__ == '__main__':image_path = "../datasets/verify/cats/13.jpeg"image = Image.open(image_path)image = image.convert("RGB")# 因为png格式是四个通道，处理RGB三通道外，还有一个透明度通道。所以，我们调用image = image.convert("RGB")，# 保留其颜色通道。当然，如果图片本来就是三颜色通道，经过此操作，不变。加上这一步后，可以适应png，jpg各种格式的图片。transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32, 32)), torchvision.transforms.ToTensor()])# module = MyModule()image = transform(image)module = torch.load("./weights/module_9.pth", map_location=torch.device('cpu'))   # 加载模型，如果使用GPU的模型用CPU来使用，需要在加载中指定CPU。image = torch.reshape(image, (1, 3, 32, 32))module.eval()with torch.no_grad():output = module(image)print(output)print(output.argmax(1))

十五、计算时间

import time'''
计算程序运行时间
'''
if __name__ == '__main__':start = time.time()  # 开始记时for i in range(2):time.sleep(1.0)end = time.time()  # 结束记时print(end - start)