Skip to content

Commit

Permalink
refactor(traffic_light_classifier): update training docs (#4917)
Browse files Browse the repository at this point in the history 
* refactor(traffic_light_classifier): update training docs

Signed-off-by: Kaan Çolak <kaancolak95@gmail.com>

---------

Signed-off-by: Kaan Çolak <kaancolak95@gmail.com>
  • Loading branch information
@kaancolak
kaancolak authored Sep 21, 2023
1 parent b87bfac commit 68e64ac
Showing 1 changed file with 224 additions and 37 deletions.
261 changes: 224 additions & 37 deletions perception/traffic_light_classifier/README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -92,67 +92,254 @@ These colors and shapes are assigned to the message as follows:
| `red_max_s` | int | the maximum saturation of red color |
| `red_max_v` | int | the maximum value (brightness) of red color |

## Customization of CNN model
## Training Traffic Light Classifier Model

### Overview

This guide provides detailed instructions on training a traffic light classifier model using the **[mmlab/mmpretrain](https://github.com/open-mmlab/mmpretrain)** repository
and deploying it using **[mmlab/mmdeploy](https://github.com/open-mmlab/mmdeploy)**.
If you wish to create a custom traffic light classifier model with your own dataset, please follow the steps outlined below.

### Data Preparation

#### Use Sample Dataset

Autoware offers a sample dataset that illustrates the training procedures for
traffic light classification. This dataset comprises 1045 images categorized
into red, green, and yellow labels. To utilize this sample dataset,
please download it from **[link](https://autoware-files.s3.us-west-2.amazonaws.com/dataset/traffic_light_sample_dataset.tar.gz)** and extract it to a designated folder of your choice.

#### Use Your Custom Dataset

To train a traffic light classifier, adopt a structured subfolder format where each
subfolder represents a distinct class. Below is an illustrative dataset structure example;

```python
DATASET_ROOT
├── TRAIN
│ ├── RED
│ │ ├── 001.png
│ │ ├── 002.png
│ │ └── ...
│ │
│ ├── GREEN
│ │ ├── 001.png
│ │ ├── 002.png
│ │ └──...
│ │
│ ├── YELLOW
│ │ ├── 001.png
│ │ ├── 002.png
│ │ └──...
│ └── ...
│
├── VAL
│ └──...
│
│
└── TEST
└── ...

Currently, in Autoware, [MobileNetV2](https://arxiv.org/abs/1801.04381v3) and [EfficientNet-b1](https://arxiv.org/abs/1905.11946v5) are provided.
The corresponding onnx files are `data/traffic_light_classifier_mobilenetv2.onnx` and `data/traffic_light_classifier_efficientNet_b1.onnx` (These files will be downloaded during the build process).
Also, you can apply the following models shown as below, for example.

- [ResNet](https://openaccess.thecvf.com/content_cvpr_2016/html/He_Deep_Residual_Learning_CVPR_2016_paper.html)
- [MobileNetV3](https://arxiv.org/abs/1905.02244)
...
```

### Installation

#### Prerequisites

**Step 1.** Download and install Miniconda from the [official website](https://mmpretrain.readthedocs.io/en/latest/get_started.html).

**Step 2.** Create a conda virtual environment and activate it

In order to train models and export onnx model, we recommend [open-mmlab/mmclassification](https://github.com/open-mmlab/mmclassification.git).
Please follow the [official document](https://mmclassification.readthedocs.io/en/latest/) to install and experiment with mmclassification. If you get into troubles, [FAQ page](https://mmclassification.readthedocs.io/en/latest/faq.html) would help you.
```bash
conda create --name tl-classifier python=3.8 -y
conda activate tl-classifier
```

The following steps are example of a quick-start.
**Step 3.** Install PyTorch

### step 0. Install [MMCV](https://github.com/open-mmlab/mmcv.git) and [MIM](https://github.com/open-mmlab/mim.git)
Please ensure you have PyTorch installed, compatible with CUDA 11.6, as it is a requirement for current Autoware

_NOTE_ : First of all, install [PyTorch](https://pytorch.org/) suitable for your CUDA version (CUDA11.6 is supported in Autoware).
```bash
conda install pytorch==1.13.1 torchvision==0.14.1 pytorch-cuda=11.6 -c pytorch -c nvidia
```

In order to install mmcv suitable for your CUDA version, install it specifying a url.
#### Install mmlab/mmpretrain

```shell
# Install mim
$ pip install -U openmim
**Step 1.** Install mmpretrain from source

# Install mmcv on a machine with CUDA11.6 and PyTorch1.13.0
$ pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu116/torch1.13/index.html
```bash
cd ~/
git clone https://github.com/open-mmlab/mmpretrain.git
cd mmpretrain
pip install -U openmim && mim install -e .
```

### step 1. Install MMClassification
### Training

MMPretrain offers a training script that is controlled through a configuration file.
Leveraging an inheritance design pattern, you can effortlessly tailor the training script
using Python files as configuration files.

You can install MMClassification as a Python package or from source.
In the example, we demonstrate the training steps on the MobileNetV2 model,
but you have the flexibility to employ alternative classification models such as
EfficientNetV2, EfficientNetV3, ResNet, and more.

```shell
# As a Python package
$ pip install mmcls
#### Create a config file

# From source
$ git clone https://github.com/open-mmlab/mmclassification.git
$ cd mmclassification
$ pip install -v -e .
Generate a configuration file for your preferred model within the `configs` folder

```bash
touch ~/mmpretrain/configs/mobilenet_v2/mobilenet-v2_8xb32_custom.py
```

### step 2. Train your model
Open the configuration file in your preferred text editor and make a copy of
the provided content. Adjust the **data_root** variable to match the path of your dataset.
You are welcome to customize the configuration parameters for the model, dataset, and scheduler to
suit your preferences

```python
# Inherit model, schedule and default_runtime from base model
_base_ = [
'../_base_/models/mobilenet_v2_1x.py',
'../_base_/schedules/imagenet_bs256_epochstep.py',
'../_base_/default_runtime.py'
]

# Set the number of classes to the model
# You can also change other model parameters here
# For detailed descriptions of model parameters, please refer to link below
# (Customize model)[https://mmpretrain.readthedocs.io/en/latest/advanced_guides/modules.html]
model = dict(head=dict(num_classes=3, topk=(1, 3)))

# Set max epochs and validation interval
train_cfg = dict(by_epoch=True, max_epochs=50, val_interval=5)

# Set optimizer and lr scheduler
optim_wrapper = dict(
optimizer=dict(type='SGD', lr=0.001, momentum=0.9))
param_scheduler = dict(type='StepLR', by_epoch=True, step_size=1, gamma=0.98)

dataset_type = 'CustomDataset'
data_root = "/PATH/OF/YOUR/DATASET"

# Customize data preprocessing and dataloader pipeline for training set
# These parameters calculated for the sample dataset
data_preprocessor = dict(
mean=[0.2888 * 256, 0.2570 * 256, 0.2329 * 256],
std=[0.2106 * 256, 0.2037 * 256, 0.1864 * 256],
num_classes=3,
to_rgb=True,
)

# Customize data preprocessing and dataloader pipeline for train set
# For detailed descriptions of data pipeline, please refer to link below
# (Customize data pipeline)[https://mmpretrain.readthedocs.io/en/latest/advanced_guides/pipeline.html]
train_pipeline = [
dict(type='LoadImageFromFile'),
dict(type='Resize', scale=224),
dict(type='RandomFlip', prob=0.5, direction='horizontal'),
dict(type='PackInputs'),
]
train_dataloader = dict(
dataset=dict(
type=dataset_type,
data_root=data_root,
ann_file='',
data_prefix='train',
with_label=True,
pipeline=train_pipeline,
),
num_workers=8,
batch_size=32,
sampler=dict(type='DefaultSampler', shuffle=True)
)

# Customize data preprocessing and dataloader pipeline for test set
test_pipeline = [
dict(type='LoadImageFromFile'),
dict(type='Resize', scale=224),
dict(type='PackInputs'),
]

# Customize data preprocessing and dataloader pipeline for validation set
val_cfg = dict()
val_dataloader = dict(
dataset=dict(
type=dataset_type,
data_root=data_root,
ann_file='',
data_prefix='val',
with_label=True,
pipeline=test_pipeline,
),
num_workers=8,
batch_size=32,
sampler=dict(type='DefaultSampler', shuffle=True)
)

val_evaluator = dict(topk=(1, 3,), type='Accuracy')

test_dataloader = val_dataloader
test_evaluator = val_evaluator

```

Train model with your experiment configuration file. For the details of config file, see [here](https://mmclassification.readthedocs.io/en/latest/tutorials/config.html).
#### Start training

```shell
# [] is optional, you can start training from pre-trained checkpoint
$ mim train mmcls YOUR_CONFIG.py [--resume-from YOUR_CHECKPOINT.pth]
```bash
cd ~/mmpretrain
python tools/train.py configs/mobilenet_v2/mobilenet-v2_8xb32_custom.py
```

### step 3. Export onnx model
Training logs and weights will be saved in the `work_dirs/mobilenet-v2_8xb32_custom` folder.

In exporting onnx, use `mmclassification/tools/deployment/pytorch2onnx.py` or [open-mmlab/mmdeploy](https://github.com/open-mmlab/mmdeploy.git).
### Convert PyTorh model to ONNX model

```shell
cd ~/mmclassification/tools/deployment
python3 pytorch2onnx.py YOUR_CONFIG.py ...
#### Install mmdeploy

The 'mmdeploy' toolset is designed for deploying your trained model onto various target devices.
With its capabilities, you can seamlessly convert PyTorch models into the ONNX format.

```bash
# Activate your conda environment
conda activate tl-classifier

# Install mmenigne and mmcv
mim install mmengine
mim install "mmcv>=2.0.0rc2"

# Install mmdeploy
pip install mmdeploy==1.2.0

# Support onnxruntime
pip install mmdeploy-runtime==1.2.0
pip install mmdeploy-runtime-gpu==1.2.0
pip install onnxruntime-gpu==1.8.1

#Clone mmdeploy repository
cd ~/
git clone -b main https://github.com/open-mmlab/mmdeploy.git
```

#### Convert PyTorch model to ONNX model

```bash
cd ~/mmdeploy

# Run deploy.py script
# deploy.py script takes 5 main arguments with these order; config file path, train config file path,
# checkpoint file path, demo image path, and work directory path
python tools/deploy.py \
~/mmdeploy/configs/mmpretrain/classification_onnxruntime_static.py\
~/mmpretrain/configs/mobilenet_v2/train_mobilenet_v2.py \
~/mmpretrain/work_dirs/train_mobilenet_v2/epoch_300.pth \
/SAMPLE/IAMGE/DIRECTORY \
--work-dir mmdeploy_model/mobilenet_v2
```

Converted ONNX model will be saved in the `mmdeploy/mmdeploy_model/mobilenet_v2` folder.

After obtaining your onnx model, update parameters defined in the launch file (e.g. `model_file_path`, `label_file_path`, `input_h`, `input_w`...).
Note that, we only support labels defined in [tier4_perception_msgs::msg::TrafficLightElement](https://github.com/tier4/tier4_autoware_msgs/blob/tier4/universe/tier4_perception_msgs/msg/traffic_light/TrafficLightElement.msg).

Expand Down

0 comments on commit 68e64ac

Please sign in to comment.