YoloV5

Here I'm gonna describe the best way I've found to run YoloV5 training and inference.

First of all my machine has an RTX 3060 so I had to do some changes to the YoloV5 dockerfile.

The following is a Dockerfile that should work with RTX 3000 series

FROM ultralytics/yolov5

RUN pip install torch==1.9.1+cu111 torchvision==0.10.1+cu111 torchaudio==0.9.1 -f https://download.pytorch.org/whl/torch_stable.html
RUN pip install roboflow

simply run the command docker build -t yolov5 . to build the docker image.

After the build is finished, to start a training do

docker run --shm-size 50G -e WANDB_MODE=disabled -v $(pwd)/shared/runs:/usr/src/app/runs --gpus all yolov5 python train.py

This should trigger the default training of the COCO small dataset.

if you want to run on custom datasets, you can use https://roboflow.com/ to create your own dataset, or download a public one.