Message Boards

GROUPS:

Staff Picks Data Science Mathematica Add-Ons Wolfram Language Machine Learning Neural Networks

Alec Graves

Tuning YOLOv2 object detection neural networks on custom datasets

Alec Graves

Posted 5 months ago

1630 Views

1 Reply

9 Total Likes

Follow this post

Hello community, please enjoy my latest project! I hope you can find it useful, and feel free to provide some feedback. I am also quite new to organizing libraries for the Wolfram Language, so any feedback there will be greatly appreciated!

Tuning YOLOv2 object detection neural networks on custom datasets
by Alec Graves

Introducing 🐺 WolfDetector 🐺

A library for fine-tuning YOLOv2 object detection neural networks on your own datasets - written in the Wolfram Language. Currently, it supports loading any PASCAL VOC style datasets for training.

Most object detection models take a lot of data to train, but we start with a pre-trained YOLOv2 model from the Wolfram Neural Net repository. In these examples, I will be training on only a few hundred images.

With this library, it is easy to build your own dataset with an open-source image labeling tool and train your own object detector for your specific domain!

If you would like to help this project grow, or view the library code, I am using GitHub to host this project’s source: https://github.com/alecGraves/WolfDetector

Demo: Fruit Detection! 🍍🍊🍓🍌

Let’s get started!

First, set your notebook directory to the location of the WolfDetector folder (containing the WolfDetector.m file).
Then import the library.

In[]:=

SetDirectory[FileNameJoin[{$HomeDirectory,"IdeaProjects","WolfDetector","WolfDetector"}]];Get["WolfDetector`"]

Now, let’s load our first dataset. I have downloaded the [Fruit Images for Object Detection dataset from Kaggle ( https://www.kaggle.com/mbkinaci/fruit-images-for-object-detection )

No modifictation to the dataset is required! The WolfDetector library knows how to read all correctly-formatted voc-style datasets.

Note that I remove the first image from the dataset so we can use it to test later!

In[]:=

dataset=WolfDetectorLoadDataset["E:\\datasets\\fruits_detection","voc"];Length[dataset]img=Import[First@Keys@dataset]dataset=Rest[dataset];Length[dataset]

Out[]=

200

Out[]=

199

We only use 199 images!

Next, we need to extract useful information about the dataset (number of classes and good bounding box priors) and convert our dataset into numeric data that will be good for training on.

In[]:=

dataInfo=WolfDetectorDataInfo[dataset]trainingDataset=WolfDetectorTrainingDataset[dataset,dataInfo];

Out[]=

classExtractor

NetEncoder



Type:	Class
Output:	vector (size: 4) of booleans



,anchors{{0.427462,0.495952},{0.0607151,0.0783495},{0.212759,0.32766},{0.156532,0.18888},{0.108254,0.15788}}

After this, we need to build our YOLOv2 neural network. Note that I use arguments to resize the network for 224x224 inputs, rather than using the default 416x416 inputs. This will speed up training.

In[]:=

{net,outputLayer}=BuildWolfDetector[dataInfo,224]loss=WolfDetectorLoss[outputLayer]

Out[]=



NetGraph



Input port:	image
Output port:	array (size: 45×7×7)



NetGraph



Input port:	array (size: 45×7×7)
Number of outputs:	3





Out[]=

NetGraph



Number of inputs:	2
Output port:	real



Next, we can begin training on our trainingDataset using our custom YOLOv2 loss function, and finally, we can run prediction.

I have enabled training on multiple GPUs with TargetDevice{“GPU”, All}, but you can remove that if you only have a CPU.

Also note that I enable gradient clipping - this helps avoid a common error “NetTrain::arrdiv: Training was stopped early because one or more trainable parameters of the net diverged.”

In[]:=

trained=NetTrain[net,trainingDataset,All,LossFunction->loss,MaxTrainingRounds->32,TargetDevice{"GPU",All},ValidationSet->Scaled[0.5],BatchSize->32,Method->{"SGD","GradientClipping"->0.001}]predictor=NetChain[{trained["TrainedNet"],outputLayer}]pred=predictor[img];outputLabels=WolfDetectorOutput[pred,dataInfo]output=WolfDetectorVisualize[img,outputLabels]

Out[]=

NetTrainResultsObject

NetTrain Results

summary

batches:

96,

rounds:

32,

time:

1.2min

examples/s:

44.

data

training examples:

96,

validation examples:

128,

processed examples:

3072,

skipped examples:

method

SGD

optimizer

batch size

32,

GPU

All

round

loss:

1.22

validation

loss:

3.6

	rounds
loss

training set

validation set



Out[]=

NetChain



Input port:	image
Number of outputs:	3



Out[]=

{Rectangle[{0.112473,0.347127},{0.666443,0.782795}]pineapple,Rectangle[{0.672314,0.76143},{0.850765,1.}]dragon fruit}

{2,2,2}

Out[]=

Note, WolfDetectorOutput can be modified to allow for less-confident predictions. In this block, I lower the confidence threshold from 0.5 to 0.4, and the less-confident snake-fruit prediction appears.

WolfDetectorOutput is currently implemented as the Non Maximum Suppression algorithm, commonly used with object detectors.

In[]:=

pred=predictor[img];outputLabels=WolfDetectorOutput[pred,dataInfo,0.4]output=WolfDetectorVisualize[img,outputLabels]

Out[]=

{Rectangle[{0.112473,0.347127},{0.666443,0.782795}]pineapple,Rectangle[{0.672314,0.76143},{0.850765,1.}]dragon fruit,Rectangle[{0.636685,0.566013},{0.733224,0.663572}]snake fruit}

{2,2,3}

Out[]=

Demo: Microcontroller Detection ⚡

Now, let’s load a new VOC-style dataset!

This time, I use the Microcontroller Detection dataset from Kaggle: https://www.kaggle.com/tannergi/microcontroller-detection - Again, no modifications required!

I remove the first image in the dataset from, as before...

In[]:=

dataset=WolfDetectorLoadDataset["E:\\datasets\\microcontroller_detection","voc"];Length[dataset]img=Import[First@Keys@dataset]dataset=Rest[dataset];Length[dataset]

Out[]=

149

Out[]=

148

Wow! Only 148 images!

Now, let’s extract this dataset’s info, build a new model, train and predict!

In[]:=

dataInfo=WolfDetectorDataInfo[dataset]trainingDataset=WolfDetectorTrainingDataset[dataset,dataInfo];{net,outputLayer}=BuildWolfDetector[dataInfo,224]loss=WolfDetectorLoss[outputLayer]trained=NetTrain[net,trainingDataset,All,LossFunction->loss,MaxTrainingRounds->32,TargetDevice{"GPU",All},ValidationSet->Scaled[0.2],BatchSize->16,LearningRate0.0002,Method->{"SGD","GradientClipping"->0.01}]predictor=NetChain[{trained["TrainedNet"],outputLayer}]pred=predictor[img];outputLabels=WolfDetectorOutput[pred,dataInfo]output=WolfDetectorVisualize[img,outputLabels]

Out[]=

classExtractor

NetEncoder



Type:	Class
Output:	vector (size: 4) of booleans



,anchors{{0.193125,0.18425},{0.523009,0.374444},{0.640761,0.647971},{0.339347,0.286023},{0.347557,0.460682}}

Out[]=



NetGraph



Input port:	image
Output port:	array (size: 45×7×7)



NetGraph



Input port:	array (size: 45×7×7)
Number of outputs:	3





Out[]=

NetGraph



Number of inputs:	2
Output port:	real



Out[]=

NetTrainResultsObject

NetTrain Results

summary

batches:

256,

rounds:

32,

time:

2.1min

examples/s:

32.

data

training examples:

128,

validation examples:

32,

processed examples:

4096,

skipped examples:

method

SGD

optimizer

batch size

16,

GPU

All

round

loss:

1.14

validation

loss:

6.62

	rounds
loss

training set

validation set



Out[]=

NetChain



Input port:	image
Number of outputs:	3



Out[]=

{Rectangle[{0.323325,0.306553},{0.694018,0.683911}]Heltec_ESP32_Lora}

{2,2,1}

Out[]=

This time, the detection is correct even with the default box confidence threshold of 0.5.

Chess ♟🐴 (Negative Example ❌)

Finally, we will try another dataset (which is not as good for training neural networks: Many of the images are not representative of the test set, and there are very few views of each chess piece in the set).

You can get this data from the Chess Piece Detection dataset on Kaggle: https://www.kaggle.com/tannergi/chess-piece-detection

This example shows the importance of taking varied pictures if you only have a few images for your domain!!

In[]:=

dataset=WolfDetectorLoadDataset["E:\\datasets\\chess_detection","voc"];Length[dataset]img=Import[First@Keys@dataset]dataset=Rest[dataset];Length[dataset]

Out[]=

In[]:=

dataInfo=WolfDetectorDataInfo[dataset]trainingDataset=WolfDetectorTrainingDataset[dataset,dataInfo];{net,outputLayer}=BuildWolfDetector[dataInfo,416]loss=WolfDetectorLoss[outputLayer]trained=NetTrain[net,trainingDataset,All,LossFunction->loss,MaxTrainingRounds->32,TargetDevice{"GPU",All},ValidationSet->Scaled[0.05],BatchSize->16,LearningRate0.0002,Method->{"SGD","GradientClipping"->0.01}]predictor=NetChain[{trained["TrainedNet"],outputLayer}]pred=predictor[img];outputLabels=WolfDetectorOutput[pred,dataInfo]output=WolfDetectorVisualize[img,outputLabels]

Out[]=

classExtractor

NetEncoder



Type:	Class
Output:	vector (size: 12) of booleans



,anchors{{0.0991906,0.19619},{0.0571086,0.113493},{0.120996,0.345033},{0.212106,0.604794},{0.172506,0.295674}}

Out[]=



NetGraph



Input port:	image
Output port:	array (size: 85×13×13)



NetGraph



Input port:	array (size: 85×13×13)
Number of outputs:	3





Out[]=

NetGraph



Number of inputs:	2
Output port:	real



Out[]=

NetTrainResultsObject

NetTrain Results

summary

batches:

128,

rounds:

32,

time:

2.0min

examples/s:

17.

data

training examples:

64,

validation examples:

16,

processed examples:

2048,

skipped examples:

method

SGD

optimizer

batch size

16,

GPU

All

round

loss:

8.28

validation

loss:

6.16×

	rounds
loss

training set

validation set



Out[]=

NetChain



Input port:	image
Number of outputs:	3



Out[]=

{Rectangle[{0.00758369,0.824639},{0.010838,0.867677}]black-pawn,Rectangle[{0.0000196616,0.676689},{0.000313856,0.861772}]black-pawn,Rectangle[{0.,0.77187},{0.0290895,0.981919}]black-pawn,Rectangle[{0.,0.350222},{0.119638,1.}]white-rook,Rectangle[{0.,0.683212},{0.0135611,1.}]white-knight,Rectangle[{0.,0.656507},{0.000224829,0.881956}]white-bishop,Rectangle[{0.,0.907341},{0.,0.951888}]white-queen,Rectangle[{0.,0.644051},{0.,0.936432}]black-bishop,Rectangle[{0.,0.476298},{0.0113501,1.}]black-king,Rectangle[{0.,0.616206},{0.0474706,1.}]black-king,Rectangle[{0.,0.822344},{0.152985,1.}]white-queen,Rectangle[{0.,0.577501},{0.013939,0.814709}]black-bishop,Rectangle[{0.0277455,0.848349},{0.136745,1.}]black-queen,Rectangle[{0.084019,0.727814},{0.10507,0.816569}]white-knight,Rectangle[{0.0666252,0.549256},{0.087851,1.}]black-pawn,Rectangle[{0.000432335,0.291727},{0.200543,1.}]black-pawn,Rectangle[{0.0523723,0.415838},{0.0985253,0.978808}]white-knight,Rectangle[{0.387019,0.359827},{0.534459,0.714312}]white-queen,Rectangle[{0.436109,0.0379277},{0.626792,0.778921}]black-bishop}

{2,2,19}

Note how the validation loss does not really decrease during training. This dataset is not very optimal for training neural networks. There are only very few images of each chess piece, and the chess pieces are not on a chess board as you would normally see them. Finally, the images are mostly in the center of the image, with a single chess piece standing alone. This kind of data is not very good for training a neural network chess piece detector, so our results are not very good.

This is why it is very important to take lots of pictures with lots of objects together, apart, and in all corners of the image, especially if there are only a hundred images. It is always important to understand your data!!

We can modify our non-max-suppression confidence threshold, but the results will not improve, since the network is trained on bad data:

In[]:=

pred=predictor[img];outputLabels=WolfDetectorOutput[pred,dataInfo,0.5,0.01]output=WolfDetectorVisualize[ImageResize[img,{416}],outputLabels]

Out[]=

{Rectangle[{0.00758369,0.824639},{0.010838,0.867677}]black-pawn,Rectangle[{0.0000196616,0.676689},{0.000313856,0.861772}]black-pawn,Rectangle[{0.,0.77187},{0.0290895,0.981919}]black-pawn,Rectangle[{0.,0.350222},{0.119638,1.}]white-rook,Rectangle[{0.,0.683212},{0.0135611,1.}]white-knight,Rectangle[{0.,0.656507},{0.000224829,0.881956}]white-bishop,Rectangle[{0.,0.907341},{0.,0.951888}]white-queen,Rectangle[{0.,0.644051},{0.,0.936432}]black-bishop,Rectangle[{0.,0.476298},{0.0113501,1.}]black-king,Rectangle[{0.,0.822344},{0.152985,1.}]white-queen,Rectangle[{0.,0.577501},{0.013939,0.814709}]black-bishop,Rectangle[{0.0277455,0.848349},{0.136745,1.}]black-queen,Rectangle[{0.084019,0.727814},{0.10507,0.816569}]white-knight,Rectangle[{0.0666252,0.549256},{0.087851,1.}]black-pawn,Rectangle[{0.387019,0.359827},{0.534459,0.714312}]white-queen,Rectangle[{0.436109,0.0379277},{0.626792,0.778921}]black-bishop}

{2,2,16}

You have reached the end of the demo!

If you want to learn more, comment below, or check out the source code here: https://github.com/alecGraves/WolfDetector

Keywords: YOLO, YOLOV2, fine tuning, pre trained, object detection, neural networks, machine learning

POSTED BY: Alec Graves

Answer

1 Reply

Sort By:

Moderation Team

Moderation Team, WOLFRAM

Posted 5 months ago

-- you have earned *Featured Contributor Badge* Your exceptional post has been selected for our editorial column *Staff Picks* http://wolfr.am/StaffPicks and Your Profile is now distinguished by a *Featured Contributor Badge* and is displayed on the Featured Contributor Board. Thank you!

POSTED BY: Moderation Team

Answer

Reply to this discussion

Reply Preview

Attachments

Remove Add a file to this post

Follow this discussion

or Discard

Group Abstract

Feedback