YOLO Training Conclusion on Custom Dataset and FAQs
Prepare dataset
You can google how to prepare your custom dataset. There are lots of tutorials available. I recommend this article.
My custom dataset has ~15k images and 3 classes. I just replace image files in VOC dataset /VOCdevkit
with my custom images and run official python script to generate data YOLO needs. You can click the official website link above to find the script. It's easy to learn how to use it.
Compile
We need to modify Makefile
before we make
it.
GPU=1 # If you have a GPU, it will speed up the training process greatly.
CUDNN=1 # Using cudnn
OPENCV=0 # Using camera
OPENMP=0
DEBUG=0
...
Then, just make
the source code.
Train
There are three files needing to be modified before training on custom dataset.
- First, you just need to write class names in your dataset to
data/voc.names
. Of course, you can change the file name.
class A
class B
class C
...
- Second, you need to modify
cfg/voc.names
. Just change it based on your custom dataset and file paths.
classes= 3
train = <path to train file>/train.txt
valid = <path to test file>/test.txt
names = data/voc.names
backup = backup
- Before strating step 3, I recommend you to read this article to understand what parameters mean.
- Third, the key step here. We need to modify yolo network configuration to meet our requirements for custom dataset. Read configuration below carefully. Here we give two examples,
yolov2-tiny-voc.cfg
andyolov3-voc.cfg
, you can configure the network you like.
// yolov2-tiny-voc.cfg
[net]
# Testing
# batch=1
# subdivisions=1
# Training
batch=64 # you need to comment testing cfg and uncomment training cfg
subdivisions=2
width=416
height=416
channels=3
...
# you can change it for your own settings
learning_rate=0.001
max_batches = 40200
policy=steps
steps=-1,100,20000,30000
scales=.1,10,.1,.1
...
...
[convolutional]
size=1
stride=1
pad=1
filters=40 # filters = num * (classes + coords + 1) = 5 * (3 + 4 + 1)
activation=linear
[region]
anchors = 1.08,1.19, 3.42,4.41, 6.63,11.38, 9.42,5.11, 16.62,10.52
bias_match=1
classes=3 # you may need to change this
coords=4 # compute filters above
num=5 # compute filters above
softmax=1
jitter=.2
rescore=1
object_scale=5
noobject_scale=1
class_scale=1
coord_scale=1
absolute=1
thresh = .6
random=1
// yolov3-voc.cfg
[net]
# Testing
# batch=1
# subdivisions=1
# Training
batch=64 # you need to comment testing cfg and uncomment training cfg
subdivisions=2
width=416
height=416
channels=3
...
# you can change it for your own settings
learning_rate=0.001
max_batches = 40200
policy=steps
steps=-1,100,20000,30000
scales=.1,10,.1,.1
...
...
[convolutional]
size=1
stride=1
pad=1
filters=24 # filters = 3 * (classes + 5)
activation=linear
[yolo]
mask = 6,7,8
anchors = 10,13, 16,30, 33,23, 30,61, 62,45, 59,119, 116,90, 156,198, 373,326
classes=3 # you need to modify
num=9
jitter=.3
ignore_thresh = .5
truth_thresh = 1
random=1
...
...
[convolutional]
size=1
stride=1
pad=1
filters=24 # filters = 3 * (classes + 5)
activation=linear
[yolo]
mask = 3,4,5
anchors = 10,13, 16,30, 33,23, 30,61, 62,45, 59,119, 116,90, 156,198, 373,326
classes=3 # you need to modify
num=9
jitter=.3
ignore_thresh = .5
truth_thresh = 1
random=1
...
...
[convolutional]
size=1
stride=1
pad=1
filters=24 # filters = 3 * (classes + 5)
activation=linear
[yolo]
mask = 0,1,2
anchors = 10,13, 16,30, 33,23, 30,61, 62,45, 59,119, 116,90, 156,198, 373,326
classes=3 # you need to modify
num=9
jitter=.3
ignore_thresh = .5
truth_thresh = 1
random=1
- Now, we can run the training command. Here we train from scratch, you also can download pre-trained weights from YOLO official website above.
./darknet detector train cfg/voc.data cfg/yolov3-voc.cfg
- In addition, you can read this article to understand outputs of YOLO.
Test
It's easy to test model trainde well. Just run command here.
./darknet detector test cfg/voc.data cfg/yolov3-voc.cfg backup/yolov3-voc_final.weights <path to testing image>
If you get no bounding box in predictions.jpg
, you firstly need to check your cfg/yolov3-voc.cfg
file. In fact, I forgot to uncomment testing configuration in .cfg
file. And just modify it, you will get correct prediction results. If nothing changes, you need to think about whether your custom dataset correctly generated.
FAQs
1. Dimension mismatch
when loading network model
Obviously, your network filters is incorrect. Go back to .cfg
file and modify filters
.
2. CUDA Out of Memory
when loading data
You need to modify cfg/yolov3-voc.cfg
. Increase subdivisions
, or decrease batch
.
3. Appearence of lots of nan
when training.
First, check your data path and data whether is correct. If that is not ok, you should focus on learning_rate
. Generally, you should decrease your learning_rate
in .cfg
file.
4. Training output obj=0.000000
It means that your network does not detect any object. Definately, it has some problems. You could tune learning_rate
and batch
in .cfg
file. Generally, you should decrease the learning_rate
or decrease the batch
.
5. No bounding boxes when testing images
I also encountered the problem. Actually, it is a careless problem. Just go back to the .cfg
file and uncomment testing configuration and comment training configuration.