best_model.pt usage - Githubissues

Chris-hughes10 / Yolov7-training

A clean, modular implementation of the Yolov7 model family, which uses the official pretrained weights, with utilities for training the model on custom (non-COCO) tasks.

GNU General Public License v3.0

116 stars 35 forks source link

best_model.pt usage #11

Closed nemeziz69 closed 1 year ago

nemeziz69 commented 1 year ago

How to properly load the best_model.pt model to run prediction on image?

Chris-hughes10 commented 1 year ago

Hi, you can do this in two ways:

using the Trainer, as described here
creating the model using model = create_yolov7_model(..., pretrained=False) then loading the state dict from the checkpoint using checkpoint = torch.load('best_model.pt', map_location="cpu"); model.load_state_dict(checkpoint["model_state_dict"])

I hope this helps!