如何將具有自定義權重的 yolov5l 模型加載到 Python 中的火炬中？

Question

我已經訓練了一個用於對象檢測和分類的 yolov5l 模型。 我想使用導出的權重來識別我正在創建的程序中的圖像。 我很難找到有關如何在 python 程序中使用 .pt 權重的任何內容。

我相信我使用了 pytorch 庫中的“torch.load”方法，但是當我嘗試： torch.load(path_to_weights)我得到一個 ModuleNotFoundError 因為沒有名為“models”的模塊。

任何幫助深表感謝。 非常感謝。

Answer 1

除了torch.load()之外，您還應該使用torch.load_state_dict()方法將經過訓練的參數加載到模型中。

您的torch.load()方法存在一些問題。 您應該將路徑參數作為字符串或 os.PathLike 對象提供。 （這些都寫在文檔中）。

我將提供一個簡單的代碼塊來向您展示方法。

#Initializing model

model = Model() # Assuming your model's name is Model

model.load_state_dict(torch.load(path_to_weights))

但不要忘記您的 path_to_weights 參數必須是字符串或 os.PathLike 對象。