Train YOLOR on Custom Data?
Object-detection technology is widely used as the backend of many applications in the industry including desktop and web applications. Also, it’s a backbone for many computer vision tasks, which include object segmentation, object tracking, object classification, object counting, etc. In the modern era, The goal of everyone regarding any application is,
“Application must be easy in use, take less processing time and provide best results”.
In previous years, many new object-detection models were released such as:
- YOLOv4/v5, YOLOR, YOLOX, and PP-YOLO
Each one has its advantages and disadvantages, but in this article, we will focus solely on YOLOR.
We will discuss the mentioned modules. All mentioned steps have been tested on Ubuntu 18.04 and 20.04 with CUDA 10.x/11.x.
- Installation of Modules
- Pretrained Object Detection
- Training on Custom Data
- Inference with Custom Weights
Installation of Modules:
Create a folder named “YOLOR”. Open terminal/cmd in the “YOLOR” folder, create a virtual environment with mentioned commands, and activate it.
python3 -m venv yolortraining #creation of virtual environment
source yolortraining/bin/activate #activation of virtualenvNote: The above step is not necessary, but recommended if you don't want to disturb Python system packages.
Now, clone the YOLO-R repository from the link, upgrade pip, and move to the cloned folder with the mentioned commands.
git clone https://github.com/WongKinYiu/yolor.git
pip install --upgrade pip
cd yolorNow, we need to install all of the libraries that will help in the training of YOLO-R. you can use the mentioned commands to install the required modules.
pip install Cython
pip install -qr requirements.txtWe now need to install some custom modules such as mish-Cuda and PyTorch-wavelets. You can use the mentioned commands to install both modules.
Mish-Cuda is an activation function that is used in training, while PyTorch-wavelets are used to compute 2D discrete wavelet etc.
#mish-cuda
git clone https://github.com/JunnYu/mish-cuda
cd mish-cuda
python setup.py build install
cd ..#pytorch-wavelets
git clone https://github.com/fbcotter/pytorch_wavelets
cd pytorch_wavelets
pip install .
cd ..
Pre-trained Object Detection:
We have installed all of the modules. We now test the detection with pre-trained weights to confirm that all of our modules are working fine. You can use the mentioned command in terminal/cmd to detect objects with the pre-trained weights.
You can now download the pre-trained weights file from this link. Move the downloaded file to the current working directory {yolor}
python3 detect.py --cfg cfg/yolor_p6.cfg --weights yolor_p6.pt --device 0 --source test.jpgIf everything is working fine, then you will be able to get results in the directory as mentioned below:
[yolor/inference/output/test.jpg]
Train YOLO-R on Custom Data:
The steps for custom training are as follows:
1- collect data
2- label data
3- split data (train & test)
4- create config files
5- start trainingStep-1: We need to create a dataset for YOLOR custom training. if you have no data, You can use the dataset from the openimages database.
YOLOR takes label data in the text(.txt) file and has the following format:
<object-class-id> <x> <y> <width> <height>
Step 2: For labeling on custom data, check out my article, Labelling data for object detection (Yolo).
Step 3: Once you have labeled your data, we now need to split our data into train and test folders. The split ratio will be dependent on the user while the normally preferable split is (80–20)%, which means 80% of the data is used for the training while 20% of the data is used for testing. *The Store images and labels with the mentioned folder architecture.
Folder structure:
├── yolor
## └── train
####└── images (folder including all training images)
####└── labels (folder including all training labels)
## └── test
####└── images (folder including all testing images)
####└── labels (folder including all testing labels)Goto yolor/utils/datasets.py(#372) and replace a filename split line with the mentioned line.
return [x.replace('images', 'labels').replace(os.path.splitext(x)[-1], '.txt') for x in img_paths]Note: Original filename split code will work fine, but sometimes if your filename is too large or has multiple dots, underscore, etc., the code will not able to get labels. Next up, we will need to change that line with the aforementioned command in that case.
Step-4: Now, we need to create the config files, which we will now use for training.
Goto {yolor/cfg/} folder, make a copy of “yolor_p6.cfg” and rename it as “yolor-custom.cfg”. Open file in some editor i.e. in Visual Studio Code. Search for the keyword “[implicit_mul]” and change the filter value inside every “[implicit_mul]” keyword, according to mentioned formula. The changing area is highlighted in Fig 3.4.
- filters = (#of classes +5) * 3
- So, if you have 1 class, then the value of the filter will be [(1+5)*3] = 18
then search for the “[yolo]” keyword, Inside every “[yolo]” search change the highlighted values mentioned in Fig 3.5, i.e. filters and #of classes.
Save the file, go to {yolor/data/} folder, make a copy of “coco.yaml” and rename it as “custom.yaml”. Open the file in some editor i.e. in Visual Studio Code. Replace all its content with the mentioned code.
train: /home/user/yolor/train/ #full path to train folder that include image and label folders
val: home/user/yolor/valid/ #full path to validation or test folder that includes image and label foldersnc: 1 # no of classes
names: ['Player'] #classesname
Save the file. Now create a file inside the “data” folder, name it “custom_classes.names” and add your classes as mentioned below.
football-player
new class #if any you havesave a file and go to the terminal.
Note: Make sure you have activated the virtual environment and your terminal path is set to the yolor folder.
Step-5: All preprocessing and configuration steps are done, it’s time to start training. open the terminal in the main “YOLOv5-cloned folder”, and run the below command.
python3 train.py --weights "" --cfg cfg/yolor-custom.cfg --data data/custom.yaml --batch-size 12 --device 0 --name yolor_custom_training --epochs 10 --img-size 448— weights= custom or pre-trained weights
— batch = batch size used in training.
— epochs = number of training epochs
— data = dataset info file
— cfg = configuration file
— img-size =image size that is use in training
If the above command will work fine, then YOLO-R training will start and you will be able to see a screen as mentioned below.
Inference with Custom Weights:
Once training is completed, You can use the mentioned command to get the results of your custom-trained model.
python3 detect.py --cfg cfg/yolor-custom.cfg --weights runs/train/yolor_custom_training/weights/best.pt --device 0 --source test.jpg --names data/custom_classes.namesOutput Sample:
That is all regarding “Train YOLOR on Custom Data”. You can try this on your data.
- Dataset creation from Videos: Article Link
- Labeling Data for Custom Training: Article Link
- Train YOLO-v5 on Custom Data: Article Link
- How to Prune and Sparse YOLOv5: Article Link
- How Hyperparameters of YOLOv5 Works: Article Link
About Me:
Muhammad Rizwan Munawar is a highly experienced professional with more than three years of work experience in Computer Vision and Software Development. He is working as a Computer Vision Engineer in Teknoir, he has knowledge and expertise in different computer vision techniques including Object Detection, Object Tracking, Pose Estimation, Object Segmentation, Segment Anything, Python, and Software Development, Embedded Systems, Nvidia Embedded Devices. In his free time, he likes to play online games and enjoys his time sharing knowledge with the community through writing articles on Medium.
Please feel free to comment below if you have any questions
