YOLOv8由pt文件中读取模型信息

Pytorch的pt模型文件中保存了许多模型信息,如模型结构、模型参数、任务类型、批次、数据集等
在先前的YOLOv8实验中,博主发现YOLOv8在预测时并不需要指定任务类型,因为这些信息便保存在pt模型中,那么,今天我们便来看看,其到底是如何加载这些参数的。

我们首先对pt文件进行一个简单介绍:

pt文格式

pt格式文件是PyTorch中用于保存张量数据的文件格式。与pth文件类似,pt文件也常用于模型的保存和加载,但更侧重于保存单个张量或一组张量数据。通过pt文件,我们可以方便地将张量数据持久化,并在需要时重新加载使用。

张量(Tensor)是PyTorch中的核心数据结构,用于表示多维数组。在深度学习中,张量常用于存储模型的参数、输入数据、中间结果等。因此,掌握pt文件的保存和加载方法对于PyTorch的使用者来说至关重要。

pt文件与pth的区别

pt.pth都是PyTorch模型文件的扩展名,但是它们的区别在于.pt文件是保存整个PyTorch模型的,而.pth文件只保存模型的参数。(其实现在似乎并没有区别了)
因此,如果要加载一个,pth文件,需要先定义模型的结构,然后再加载参数;而如果要加载一个,pt文件,则可以直接加载整个模型。

如何保存pt格式文件

PyTorch中,我们可以使用torch.save()函数将张量数据保存到pt文件中。

下面是一个简单的示例:

import torch
# 创建一个张量
tensor = torch.tensor([[1, 2, 3], [4, 5, 6]])
# 将张量保存到pt文件中
torch.save(tensor, 'tensor.pt')

在上面的代码中,我们首先创建了一个二维张量tensor,然后使用torch.save()函数将其保存到名为tensor.pt的文件中。保存的文件将包含张量的数据和元数据,以便在加载时能够准确地恢复张量的结构和内容。

除了保存单个张量外,我们还可以保存多个张量到一个pt文件中。这可以通过将多个张量放入一个字典或列表中,然后将整个字典或列表保存到文件中实现。

例如:

# 创建多个张量
tensor1 = torch.tensor([1, 2, 3])
tensor2 = torch.tensor([[4, 5], [6, 7]])

# 将张量放入字典中
tensors_dict = {'tensor1': tensor1, 'tensor2': tensor2}

# 将字典保存到pt文件中
torch.save(tensors_dict, 'tensors_dict.pt')

如何加载pt格式文件

加载pt文件同样使用torch.load()函数。

下面是一个加载pt文件的示例:

# 加载单个张量的pt文件
loaded_tensor = torch.load('tensor.pt')
print(loaded_tensor)

# 加载包含多个张量的字典的pt文件
loaded_dict = torch.load('tensors_dict.pt')
print(loaded_dict['tensor1'])
print(loaded_dict['tensor2'])

在加载单个张量的pt文件时,我们直接调用torch.load()函数并传入文件名即可。加载得到的loaded_tensor将是一个与原始张量结构和内容相同的张量对象。
当加载包含多个张量的字典的pt文件时,我们同样使用torch.load()函数。加载得到的loaded_dict将是一个字典对象,其中包含了我们在保存时放入的所有张量。我们可以通过字典的键来访问这些张量。

强烈建议只保存模型参数,而非保存整个网络。PyTorch 官方也是这么建议的。

torch.save(net.state_dict(),path2)#只保留模型参数

(只保存模型参数)是官方推荐的方法,运行速度快,且占空间较小。需要注意的是 net.state_dict() 是将网络参数保存为字典形式(OrderedDict)load_state_dict() 加载的并不是网络参数的pth文件,而是字典。

pt文件保存神经网络

在评估时,记住一定要使用model.eval()来固定dropout和归一化层,否则每次推理会生成不同的结果。

import torch, glob, cv2
from torchvision import transforms
import numpy as np
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):  # 神经网络部分用你自己的
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 2, 1)  # nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)
        self.conv2 = nn.Conv2d(32, 64, 3, 2, 1)
        self.conv3 = nn.Conv2d(64, 128, 3, 1)
        self.dropout1 = nn.Dropout2d(0.25)
        self.dropout2 = nn.Dropout2d(0.5)
        self.fc1 = nn.Linear(6272, 128)  # 6272=128*7*7
        self.fc2 = nn.Linear(128, 8)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = self.conv3(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        self.output = F.log_softmax(x, dim=1)
        out1 = x
        return self.output,out1

def predict_mine():
    model=Net()
    model.load_state_dict(torch.load("model.pt"))
    print(model)
    images=torch.rand((1,1,64,64))
    x=model(images)
    print(x)
def torch_script_save():
    model=Net()
if __name__ == '__main__':
    save_model()
    predict_mine()

在这里插入图片描述

可以看到,我们可以通过pt文件读取出来下面的信息:

在这里插入图片描述

同时,我们也看到,我们虽然可以使用pt文件保存模型结构,但我们在推理时,依旧需要我们能够生成Net对象才能加载其数据,这其实很不方便,那么,有什么办法可以真正的将模型结构保存进去,让我们在推理过程中不需要再定义相关的类与对象呢,先前博主所使用的ONNX便是其中的一种,但它其实是另一种文件结构了,pt文件真的就不能摆脱环境吗,答案是否定的,TorchScript模型便解决了这个问题。

TorchScript模型

事实上,PyTorch提供了两种主要的模型保存和加载机制,一种是基于Python的序列化,另一种是TorchScript

普通的PyTorch模型(基于Python的序列化):

  • 保存: 使用torch.save(model.state_dict(), 'model_path.pth'),它保存了模型的权重和参数,但不保存模型的结构。(当然也是可以保存的,但我们需要处理一下才能用,比如定义好Net类)
  • 加载: 首先,您需要有模型的类定义。 创建该类的一个实例。 使用model.load_state_dict(torch.load('model_path.pth'))来加载权重。
    特点:
  • 需要Python环境和模型的原始代码来加载和运行模型。 保存的文件是Python特定的,并且依赖于特定的类结构。 主要用于继续训练或在Python环境中进行推断。

TorchScript模型:

TorchScript是PyTorch的一个子集,它创建了一个可以独立于Python运行的序列化模型。 生成方法:

  • Tracing: 使用torch.jit.trace方法。这涉及到通过模型运行一个输入示例,从而跟踪模型的执行路径。
  • Scripting: 使用torch.jit.script方法。这转化Python代码到TorchScript,允许更复杂的模型和控制流。
  • 保存: 使用torch.jit.save(traced_model, 'model_path.pt')
  • 加载: 使用torch.jit.load(‘model_path.pt’)。注意,加载不需要原始的模型类定义。

特点:

  • 可以在没有Python运行时的环境中运行,如C++
  • 提供了一种方法,将模型从Python转移到其他平台或部署环境。
  • 包含模型的完整定义,包括结构、权重和参数。

Tracing方法:

example_input = torch.randn(1, 10)
traced_model = torch.jit.trace(model, example_input)
torch.jit.save(traced_model, 'traced_model.pt')

Scripting方法:

scripted_model = torch.jit.script(model)
torch.jit.save(scripted_model, 'scripted_model.pt')

加载模型:

loaded_model = torch.jit.load('model_path.pt')

例程:

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 10)

    def forward(self, x):
        return self.fc(x)

model = SimpleModel()

# Tracing
example_input = torch.randn(1, 10)
traced_model = torch.jit.trace(model, example_input)
torch.jit.save(traced_model, 'traced_simple_model.pt')

# Scripting
scripted_model = torch.jit.script(model)
torch.jit.save(scripted_model, 'scripted_simple_model.pt')

# 加载模型
loaded_model = torch.jit.load('traced_simple_model.pt')

我们采用TorchScript结构去执行先前的Net

def torch_script_save():
    model=Net()
    example_input =torch.rand((1,1,64,64))
    traced_model = torch.jit.trace(model, example_input)
    torch.jit.save(traced_model, 'traced_simple_model.pt')

    # Scripting
    scripted_model = torch.jit.script(model)
    torch.jit.save(scripted_model, 'scripted_simple_model.pt')
def predict_script():
    model1=torch.jit.load("traced_simple_model.pt")
    image =torch.rand((1,1,64,64))
    print(model1)
    model1.eval()
    x=model1(image)
    print(x)
    model2=torch.jit.load("scripted_simple_model.pt")
    image =torch.rand((1,1,64,64))
    print(model2)
    model2.eval()
    x=model2(image)
    print(x)

在这里插入图片描述

yaml文件内容如下:

{'nc': 1000, 
'scales': {'n': [0.33, 0.25, 1024], 's': [0.33, 0.5, 1024], 'm': [0.67, 0.75, 1024], 'l': [1.0, 1.0, 1024], 'x': [1.0, 1.25, 1024]}, 
'backbone': [[-1, 1, 'Conv', [64, 3, 2]], [-1, 1, 'Conv', [128, 3, 2]], [-1, 3, 'C2f', [128, True]], [-1, 1, 'Conv', [256, 3, 2]], [-1, 6, 'C2f', [256, True]], [-1, 1, 'Conv', [512, 3, 2]], [-1, 6, 'C2f', [512, True]], [-1, 1, 'Conv', [1024, 3, 2]], [-1, 3, 'C2f', [1024, True]]],
 'head': [[-1, 1, 'Classify', ['nc']]], 
 'scale': 'n',
 'yaml_file': 'yolov8n-cls.yaml', 
 'ch': 3}
  • 14
    点赞
  • 15
    收藏
    觉得还不错? 一键收藏
  • 打赏
    打赏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

彭祥.

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值