HybridNets: End-to-End Perception Network
by Dat Vu, Bao Ngo, Hung Phan π§ FPT University
(π§) corresponding author.
arXiv technical report (arXiv 2203.09035)
Table of Contents
HybridNets is an end2end perception network for multi-tasks. Our work focused on traffic object detection, drivable area segmentation and lane detection. HybridNets can run real-time on embedded systems, and obtains SOTA Object Detection, Lane Detection on BDD100K Dataset.
HybridNets
β backbone.py # Model configuration
| export.py # UPDATED 10/2022: onnx weight with accompanying .npy anchors
β hubconf.py # Pytorch Hub entrypoint
β hybridnets_test.py # Image inference
β hybridnets_test_videos.py # Video inference
β train.py # Train script
β train_ddp.py # DistributedDataParallel training (Multi GPUs)
β val.py # Validate script
β val_ddp.py # DistributedDataParralel validating (Multi GPUs)
β
ββββencoders # https://github.com/qubvel/segmentation_models.pytorch/tree/master/segmentation_models_pytorch/encoders
β ...
β
ββββhybridnets
β autoanchor.py # Generate new anchors by k-means
β dataset.py # BDD100K dataset
β loss.py # Focal, tversky (dice)
β model.py # Model blocks
β
ββββprojects
β bdd100k.yml # Project configuration
β
ββββros # C++ ROS Package for path planning
β ...
β
ββββutils
| constants.py
β plot.py # Draw bounding box
β smp_metrics.py # https://github.com/qubvel/segmentation_models.pytorch/blob/master/segmentation_models_pytorch/metrics/functional.py
β utils.py # Various helper functions (preprocess, postprocess, eval...)
The project was developed with Python>=3.7 and Pytorch>=1.10.
git clone https://github.com/datvuthanh/HybridNets
cd HybridNets
pip install -r requirements.txt
# Download end-to-end weights
curl --create-dirs -L -o weights/hybridnets.pth https://github.com/datvuthanh/HybridNets/releases/download/v1.0/hybridnets.pth
# Image inference
python hybridnets_test.py -w weights/hybridnets.pth --source demo/image --output demo_result --imshow False --imwrite True
# Video inference
python hybridnets_test_videos.py -w weights/hybridnets.pth --source demo/video --output demo_result
# Result is saved in a new folder called demo_result
Recommended dataset structure:
HybridNets
ββββdatasets
ββββimgs
β ββββtrain
β ββββval
ββββdet_annot
β ββββtrain
β ββββval
ββββda_seg_annot
β ββββtrain
β ββββval
ββββll_seg_annot
ββββtrain
ββββval
Update your dataset paths in projects/your_project_name.yml
.
For BDD100K:
1) Edit or create a new project configuration, using bdd100k.yml as a template. Augmentation params are here.
# mean and std of dataset in RGB order
mean: [0.485, 0.456, 0.406]
std: [0.229, 0.224, 0.225]
# bdd100k anchors
anchors_scales: '[2**0, 2**0.70, 2**1.32]'
anchors_ratios: '[(0.62, 1.58), (1.0, 1.0), (1.58, 0.62)]'
# BDD100K officially supports 10 classes
# obj_list: ['person', 'rider', 'car', 'truck', 'bus', 'train', 'motorcycle', 'bicycle', 'traffic light', 'traffic sign']
obj_list: ['car']
obj_combine: ['car', 'bus', 'truck', 'train'] # if single class, combine these classes into 1 single class in obj_list
# leave as empty list ([]) to not combine classes
seg_list: ['road',
'lane']
seg_multilabel: false # a pixel can belong to multiple labels (i.e. lane line + underlying road)
dataset:
dataroot: path/to/imgs
labelroot: path/to/det_annot
segroot:
# must be in correct order with seg_list
- path/to/da_seg_annot
- path/to/ll_seg_annot
fliplr: 0.5
flipud: 0.0
hsv_h: 0.015
hsv_s: 0.7
hsv_v: 0.4
...
python train.py -p bdd100k -c 3 -n 16 -b 16 --lr 1e-5 --num_epochs 200
python train.py -p bdd100k # your_project_name
-c 3 # coefficient of effnet backbone, result from paper is 3
OR -bb repvgg_b0 # change your backbone with timm
-n 4 # num_workers
-b 8 # batch_size per gpu
-w path/to/weight # use 'last' to resume training from previous session
--freeze_det # freeze detection head, others: --freeze_backbone, --freeze_seg
--lr 1e-5 # learning rate
--optim adamw # adamw | sgd
--num_epochs 200
Please check python train.py --help
for cheat codes.
IMPORTANT (deprecated): If you want to train on multiple gpus, use train_ddp.py
. Tested on NVIDIA DGX with 8xA100 40GB.
Why didn't we combine DDP into the already existing train.py
script?
- Lots of if-else.
- Don't want to break functioning stuffs.
- Lazy.
Update 24/06/2022: train_ddp.py
broke because we have a lot of things changed. Therefore, we decided to write a merged train.py
with DDP support for easier maintainance. In the meantime, please clone this commit with a working train_ddp.py
script if you really have to.
python val.py -w checkpoints/weight.pth
Again, check python val.py --help
for god mode.
Validation process got killed! What do I do? => This is because we use a default confidence threshold of 0.001 to compare with other networks. So when calculating metrics, it has to handle a large amount of bounding boxes, leading to out-of-memory, and finally exploding the program before the next epoch.
That being said, there are multiple ways to circumvent this problem, choose the best that suit you:
- Train on a high-RAM instance (RAM as in main memory, not VRAM in GPU). For your reference, we can only val the combined
car
class with 64GB RAM. - Train with
python train.py --cal_map False
to not calculate metrics when validating. This option will only print validation losses. When the losses seem to flatten and the weather is nice, rent a high-RAM instance to validate the best weight withpython val.py -w checkpoints/xxx_best.pth
. We actually did this to save on cost. - Reduce the confidence threshold with
python train.py --conf_thres 0.5
orpython val.py --conf_thres 0.5
, depending on your application and end goals. You don't have to get best recall unless you're either helping us by experimenting π₯° or competing with us π .
python export.py -w checkpoints/weight.pth --width 640 --height 384
This automatically creates an ONNX weight and an anchor_{H}_{W}.npy
file to use in postprocessing. Refer to the ROS section for usage example.
If your dataset is intrinsically different from COCO or BDD100K, or the metrics of detection after training are not as high as expected, you could try enabling autoanchor in project.yml
:
...
model:
image_size:
- 640
- 384
need_autoanchor: true # set to true to run autoanchor
pin_memory: false
...
This automatically finds the best combination of anchor scales and anchor ratios for your dataset. Then you can manually edit them project.yml
and disable autoanchor.
If you're feeling lucky, maybe mess around with base_anchor_scale in backbone.py
:
class HybridNetsBackbone(nn.Module):
...
self.pyramid_levels = [5, 5, 5, 5, 5, 5, 5, 5, 6]
self.anchor_scale = [1.25,1.25,1.25,1.25,1.25,1.25,1.25,1.25,1.25,]
self.aspect_ratios = kwargs.get('ratios', [(1.0, 1.0), (1.4, 0.7), (0.7, 1.4)])
...
and model.py
:
class Anchors(nn.Module):
...
for scale, ratio in itertools.product(self.scales, self.ratios):
base_anchor_size = self.anchor_scale * stride * scale
anchor_size_x_2 = base_anchor_size * ratio[0] / 2.0
anchor_size_y_2 = base_anchor_size * ratio[1] / 2.0
...
to get a grasp on how anchor boxes work.
And because a picture is worth a thousand words, you can visualize your anchor boxes in Anchor Computation Tool.
We experimented with training stages and found that this settings achieved the best results:
--freeze_seg True
~ 200 epochs--freeze_backbone True --freeze_det True
~ 50 epochs- Train end-to-end ~ 50 epochs
The reason being detection head is harder to converge early on, so we basically skipped segmentation head to focus on detection first.
Result | Visualization | |||||||||||||||||||||
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
Result | Visualization | ||||||||||||
---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
Result | Visualization | ||||||||||||||||||
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
Original footage courtesy of Hanoi Life
Distributed under the MIT License. See LICENSE
for more information.
Our work would not be complete without the wonderful work of the following authors:
If you find our paper and code useful for your research, please consider giving a star β and citation π :
@misc{vu2022hybridnets,
title={HybridNets: End-to-End Perception Network},
author={Dat Vu and Bao Ngo and Hung Phan},
year={2022},
eprint={2203.09035},
archivePrefix={arXiv},
primaryClass={cs.CV}
}