基于Pytorch的模型推理

训练部分说明

假设我们现在有两个文件
{
first_file: train.py #用于训练模型
second_file: inference.py#用于推理检测
}
在train.py文件中我们使用了定义了一个类，里面声明了我的网络模型，例如。

class Net(nn.Moudle):
........

假设在train.py文件中，我们处理图片是用以下的代码

transforms = transforms.Compose([transforms.Resize((224, 224)),
                                            transforms.ToTensor(),
                                            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]
                                          )

然后我们最终保存权重文件是通过以下方式。

torch.save(model.state_dict,save_path)#save_path是自定义的路径，注意要加上自定义的文件名。
如：D:/my_file/first_moedl.pth  #其中first_model.pth是保存之后的名字

推理部分实现

我们需要优先导入网络框架，也就是上文提到的Net类

from train import Net

然后还有一些常用的功能包

from PIL import Image
import torch
import torch.nn as nn
from torchvision import transforms

接着便开始代码的编写了

def prediect(device):
	net = Net()#实例化net
    net.load_state_dict(torch.load('D:/my_file/first_model.pth'))#加载模型
    net = net.to(device)#同样用GPU
    torch.no_grad()
    for i in range(1000):
        img = Image.open("存放图片路径")
        image = transforms(img).unsqueeze(0)#由于训练的时候还有一个参数，是batch_size,而推理的时候没有，所以我们为了保持维度统一，就得使用.unsqueeze(0)来拓展维度
        image = image.to(device)#同样将图片数据放入cuda(GPU)中
        output = net(image)
        _, pre = torch.max(output, 1)#拿出最高置信度的结果

        print(pre) #打印出结果

if __name__ == '__main__':
	#图片格式的转化，和训练时一样
	transforms = transforms.Compose([transforms.Resize((15, 15)),
	                                            transforms.ToTensor(),
	                                            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]
	                                          )
	#我训练的时候用的是GPU，所以这里也一样
    device = torch.device("cuda")
    #开始预测输出
    prediect(device)

当然，在很多时候我们并不会把训练和推理放在两个文件夹，也不会直接打印出它的结果，也是在一个文件夹中实现了训练、测试、推理三个部分，其中我们通常会把推理之后的结果比上正确的结果，来查看准确率。但实现过程都是换汤不换药的。