Where do you get the weights used for transfer learning?

user_14863794 0
Viewed 543

I'm building a model that detects wild boars and deer to build a animal classifier.( i have only wild boar and deer dataset ).

I have to use a model in yolov5. I want to get a good model through transfer learning, but I don't know how to get the weights(pretrained) to classify wild boar and deer. how do i get it?

Or i just download yolov5s.pt and i add --weights yolov5s.pt code when training the model?

i am a beginner if letting me know what I'm doing wrong Thanks

yolov5 torch yolo deep-learning
1 Answers

Basically you want a line like this to train a new model starting from the default yolo model.

!python train.py --img 960 --batch 16 --epochs 10 --data yolo_data.yaml --weights yolov5s.pt --cache --exist-ok

When that completes you will find the weights in this folder: runs/train/exp/weights/best.pt

Then you can use that model to predict more photos like this:

model = torch.hub.load('ultralytics/yolov5', 'custom', path='runs/train/exp/weights/best.pt', force_reload=True) 
imgs = ['0001.jpeg']  # batch of images
results = model(imgs)

I have a sample notebook that may help you here https://github.com/pylabel-project/samples/blob/main/yolov5_training.ipynb

This sample uses a dataset with 2 classes: squirrels and nuts

user_51837350
Related