如何使用pt文件

Question

我正在尝试建立一个货币识别模型，我使用了kaggle上的数据集和使用yolov5的colab，并且我准确地执行了在yolov5 github上解释的步骤。最后，我下载了一个具有模型权重的.pt文件，现在我想在python文件中使用它来检测和识别货币。怎么做？

我是一个计算机视觉初学者，我完全不知道该怎么做。我在一遍又一遍地寻找，但我什么也没有触及。

import torch

# Model
model=torch.load('E:\_best.pt')

# Images
imgs=['E:\Study\currency.jpg']

# Inference
results = model(imgs)

# Results
results.print()
results.save()  # or .show()
results.show()
results.xyxy[0]  # img1 predictions (tensor)
results.pandas().xyxy[0]

Answer 1

如果希望从.pt文件中读取经过训练的参数并将其加载到模型中，则可以执行以下操作。

file = "model.pt"
model = your_model()
model.load_state_dict(torch.load(file))
# this will automatically load the file and load the parameters into the model.

在调用load_state_dict()之前，请确保.pt文件只包含模型参数，否则会发生错误。这可以通过打印(torch.load(文件))进行检查。