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