Search code examples
deep-learningtorchyoloyolov5

Where do you get the weights used for transfer learning?


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


Solution

  • 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