Search code examples
pythonyoloyolov5

Downsize image when training yolov5 model


I am looking into making my custom YoloV5 model faster, based on my current results where I have trained a on ~20k (1280 × 960) images with a configuration based on yolov5l6.yaml (A P6 model I assume)

This gives me a very well performing model, and I would now like to explore smaller/simpler models.

Option A: Is the simple and try to train with a medium/small/nano version of the P6 model

Option B: I would like to try going to an image size of 640 by using a P5 model.

My question is: do I need to manually resize my images to 640x640 to efficiently train a P5 model (e.g. based on yolov5m.yaml), or can I simply pass the desired image size as a parameter to train.py?


Solution

  • The imgsz flag determines the size of the images fed to the model, this is also true for training. So you can simply do the following

    python train.py ... --imgsz 640  # training on 640 images
    

    To check that this is actually the case you can

    print(img.shape)
    

    just before the image batch is fed to the model

    Reference