I have trained Yolov5 on a dataset(1) with X, Y, and Z classes and got weights file (A) using Yolov5s.pt weights. The model is not detecting well on Y class instances because Ys are not many in the dataset.
Now, I have a different separated dataset (2) with Y and W classes. I want to train dataset(2) and also want to save weights in weight file(A). My goal is to get one weight file that detects all classes (WXYZ). What would be the best approach to it?
4 Likes
@khandriod this is a common problem, mixing and matching datasets. The good news is you can train on multiple datasets simultaneously with YOLOv5 just by adding them to your train, val, test fields in your data.yaml.
The bad news is the classes must all correspond for this to work correctly, i.e. you can train on COCO plus another dataset at the same time, but since COCO classes are labelled 0-79, your new dataset class labels must start at 80 onward.
See GlobalWheat as an example of training multiple directories togethor:
3 Likes
@khandriod to explain a bit further, in the future we want to make this way easier, but for now sure you can train on multiple datasets super easily, the only issue is that the label classes must play well with each other. If you have DATASET A with classes APPLE, ORANGE, and you have DATASET B with BANANA, you can:
- Train two models independently, one on each dataset: (APPLE 0, ORANGE 1 in dataset 1), (BANANA 0 in dataset 2)
- Train one model on the combined datasets making sure each class index is unique: APPLE 0, ORANGE 1, BANANA 2
2 Likes
To add to what @Glenn has posted, an existing weight file cannot have new classes added to it. A new weight file will need to be trained with all the required classes and datasets.
@khandriod The options in your specific case are
- Combine your two datasets and train a new model.
- Train a new model with multiple datasets as explained by @Glenn
Whatever route you decide to go, you will need to train a new model if you want it to detect all of the classes.
2 Likes
@Glenn Thanks for the detailed reply. If I go with the first option and train Model-1 with Dataset A with classes APPLE AND ORANGES with this command. Emphasis on [-- weights option]
python train.py --img 640 --batch 16 --epochs 3 --data data_a.yaml --weights yolov5s.pt
From this command, I would get the best_a.pt weights file.
- For Model-2 with Dataset B class BANANA.
python train.py --img 640 --batch 16 --epochs 3 --data data_b.yaml --weights best_a.pt
This command will give me the best_b.pt
Now if I run a test with best_b.pt on an image would it detect all three classes (Apple, Orange, and Banana)? In other words, what weight file would be needed for training Model-2?
Thanks
2 Likes
@khandroid Please refer to my post above. In the example you gave best_b.pt will only detect the class Banana. You need to train a new model that accepts both datasets and outputs for example best_c.pt
To achieve your desired result of a model that detects the three classes (Apple, Orange and Banana) a model has to be trained with all three classes and both datasets. The easiest way to achieve this is to create a modified YAML that points to the images in both datasets.
For example copy data_a.yaml to data_c.yaml and modify it to include all images as below:
train: # train images (relative to 'path')
- images/data_a
- images/data_b
...
nc: 3
names: [
'apple',
'orange'
'banana'
]
Then train a new model on this new YAML:
python train.py --img 640 --batch 16 --epochs 3 --data data_c.yaml --weights yolov5s.pt
The final model will now detect all classes specified in data_c.yaml.
**Be sure that your labels match the index in the data_c.yaml. If data_b has banana set as index 0 you will need to update all the labels and change the index to 2.
2 Likes
@KalenMichael Thanks for the explanation. I understand your point. One more question about the index in .yaml file.
Say Dataset_A have [Apple, Oranges, Banana] classes with index (0,1,2);
and Dataset_B have [Banana, Mango] class with index (0,1)
what would be the index of classes in combine.yaml?
1 Like
@khandriod Provided that the Banana classes in Dataset_A and Dataset_B are the same you can combine them to be Apple[0], Oranges[1], Banana[2], Mango[3].
What is important is that the labels have the correct index. You should have a corresponding label file for each image file. Each label file will have a row for each object as so:
class x_center y_center width height
The ‘class’ should be the correct index as set in the names array in the YAML. So what you will need to do is alter the labels in Dataset_B by looping through all files and each row and updating ‘class’ 0 for 2 (banana) and ‘class’ 1 for 3 (mango). This will combine the classes.
As you are changing the labels remember to backup your original Dataset_B in case you need it in the future.
If you are interested to read more about creating custom datasets you can check out the wiki:
Create Custom Dataset
5 Likes
@Glenn I have a question, is it possible to train one model on two datasets one to classify objects and the second is to locate the object in the image?
@myasser63 I don’t understand your question. YOLOv5 already does the tasks of ‘classify objects and the second is to locate the object in the image’.
I mean that I have two datasets of the same objects but one dataset is divided into classes for each class ex: apple, banana, orange, pineapple while the other dataset of the same classes but annotated for detection of the object group ex Fruits. So i want to train one model on the 1st dataset is to classify classes accuratly while the other dataset is more accurate on training the model for localization.
I just wanted to know,
Query 1: I want to combine coco dataset’s 80 classes that are already trained, with 1 more class of mine (say currency), so do I have to retrain the model for all these 81 classes? (using retraining on coco and 1 new class?) Or can I combine the weights of coco and 1 new class.
Query 2: If I have to retrain on coco as well, won’t it take too long?
@kashishnaqvi101 yes, training on multiple datasets will require more time than training on a single dataset.
Can you just tell me, can I put two different directory paths one for coco and one for my custom dataset?
would this work for multiple datasets training ?
I have not added path variable but have simply given it entirely for coco.
For my custom dataset, just have given its paths till the images
train:
-datasets/coco128/images/train2017
-kashish/images/train/
val:
-datasets/coco128/images/train2017
-kashish/images/val/
@kashishnaqvi101 your YAML syntax is incorrect. Lists need a space after each dash. See YAML Syntax — Ansible Documentation
Used this
download: https://ultralytics.com/assets/coco128.zip
train:
- coco128/images/train2017
- kashish/images/train/
val:
- coco128/images/train2017
# Classes
nc: 81 # number of classes
names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
'hair drier', 'toothbrush', 'money'] # class names
But got this? why? I also see a coco128 folder above yolov5 folder now
Traceback (most recent call last):
File "/content/yolov5/utils/datasets.py", line 410, in __init__
raise Exception(f'{prefix}{p} does not exist')
Exception: train: /content/yolov5/coco128/images/train2017 does not exist
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "train.py", line 643, in <module>
main(opt)
File "train.py", line 539, in main
train(opt.hyp, opt, device, callbacks)
File "train.py", line 227, in train
prefix=colorstr('train: '), shuffle=True)
File "/content/yolov5/utils/datasets.py", line 110, in create_dataloader
prefix=prefix)
File "/content/yolov5/utils/datasets.py", line 415, in __init__
raise Exception(f'{prefix}Error loading data from {path}: {e}\nSee {HELP_URL}')
Exception: train: Error loading data from ['/content/yolov5/coco128/images/train2017', '/content/yolov5/kashish/images/train']: train: /content/yolov5/coco128/images/train2017 does not exist
See https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data
[ ]
I am getting this error, I have even put coco in my folder, still no help
Traceback (most recent call last):
File "train.py", line 643, in <module>
main(opt)
File "train.py", line 539, in main
train(opt.hyp, opt, device, callbacks)
File "train.py", line 227, in train
prefix=colorstr('train: '), shuffle=True)
File "/content/yolov5/utils/datasets.py", line 110, in create_dataloader
prefix=prefix)
File "/content/yolov5/utils/datasets.py", line 415, in __init__
raise Exception(f'{prefix}Error loading data from {path}: {e}\nSee {HELP_URL}')
Exception: train: Error loading data from ['/content/kashish/images/coco/train2017', '/content/kashish/images/train']: train: /content/kashish/images/train does not exist
See https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data
My yaml file is
path: ../kashish/
train:
- images/coco/train2017
- images/train
val:
- images/val
# Classes
nc: 81 # number of classes
names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
'hair drier', 'toothbrush', 'money'] # class names
@kashishnaqvi101 To train correctly your data must be in YOLOv5 format. Please see our Train Custom Data tutorial for full documentation on dataset setup and all steps required to start training your first model. A few excerpts from the tutorial:
1.1 Create dataset.yaml
COCO128 is an example small tutorial dataset composed of the first 128 images in COCO train2017. These same 128 images are used for both training and validation to verify our training pipeline is capable of overfitting. data/coco128.yaml, shown below, is the dataset config file that defines 1) the dataset root directory path and relative paths to train / val / test image directories (or *.txt files with image paths), 2) the number of classes nc and 3) a list of class names:
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../datasets/coco128 # dataset root dir
train: images/train2017 # train images (relative to 'path') 128 images
val: images/train2017 # val images (relative to 'path') 128 images
test: # test images (optional)
# Classes
nc: 80 # number of classes
names: [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
'hair drier', 'toothbrush' ] # class names
1.2 Create Labels
After using a tool like Roboflow Annotate to label your images, export your labels to YOLO format, with one *.txt file per image (if no objects in image, no *.txt file is required). The *.txt file specifications are:
- One row per object
- Each row is
class x_center y_center width heightformat. - Box coordinates must be in normalized xywh format (from 0 - 1). If your boxes are in pixels, divide
x_centerandwidthby image width, andy_centerandheightby image height. - Class numbers are zero-indexed (start from 0).
The label file corresponding to the above image contains 2 persons (class 0) and a tie (class 27):
1.3 Organize Directories
Organize your train and val images and labels according to the example below. YOLOv5 assumes /coco128 is inside a /datasets directory next to the /yolov5 directory. YOLOv5 locates labels automatically for each image by replacing the last instance of /images/ in each image path with /labels/. For example:
../datasets/coco128/images/im0.jpg # image
../datasets/coco128/labels/im0.txt # label
Good luck 
and let us know if you have any other questions!