mirror of https://github.com/alibaba/EasyCV.git
152 lines
4.2 KiB
Markdown
152 lines
4.2 KiB
Markdown
# YOLOX Compression Tutorial
|
|
|
|
## ENV preparation
|
|
### install torch
|
|
blade_compression needs torch >= 1.10.0
|
|
```shell
|
|
conda install pytorch==1.10.0 torchvision==0.11.0 torchaudio==0.10.0 cudatoolkit=10.2 -c pytorch
|
|
```
|
|
### install mmcv
|
|
with torch >=1.10.0, mmcv >=1.3.16 is required
|
|
```shell
|
|
pip install mmcv-full==1.4.4 -f https://download.openmmlab.com/mmcv/dist/{cu_version}/{torch_version}/index.html
|
|
```
|
|
### install onnx and pai_nni
|
|
```shell
|
|
pip install onnx
|
|
pip install http://pai-nni.oss-cn-zhangjiakou.aliyuncs.com/release/2.6.1/pai_nni-2.6.1-py3-none-manylinux1_x86_64.whl
|
|
```
|
|
### install blade_compression
|
|
```shell
|
|
pip install http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/third_party/blade_compression-0.0.2-py3-none-any.whl
|
|
```
|
|
|
|
## Data preparation
|
|
To download the dataset, please refer to [prepare_data.md](../prepare_data.md).
|
|
|
|
The data is used to train and test compression model. So the data format should be the same of train data format.
|
|
|
|
### COCO format
|
|
To use coco data to eval model, you can refer to [configs/detection/yolox/yolox_s_8xb16_300e_coco.py](../../configs/detection/yolox/yolox_s_8xb16_300e_coco.py) for more configuration details.
|
|
|
|
### PAI-Itag detection format
|
|
To use pai-itag detection format data to eval detection, you can refer to [configs/detection/yolox/yolox_s_8xb16_300e_coco_pai.py](../../configs/detection/yolox/yolox_s_8xb16_300e_coco_pai.py) for more configuration details.
|
|
|
|
## Local & PAI-DSW
|
|
|
|
To use COCO format data, use config file `configs/detection/yolox/yolox_s_8xb16_300e_coco.py`
|
|
|
|
To use PAI-Itag format data, use config file `configs/detection/yolox/yolox_s_8xb16_300e_coco_pai.py`
|
|
|
|
|
|
### Compression
|
|
**Quantize:**
|
|
This is used to quantize yolox model; The quantized model will be saved in work_dir.
|
|
|
|
```shell
|
|
python tools/quantize.py \
|
|
${CONFIG_PATH} \
|
|
${MODEL_PATH} \
|
|
--work_dir ${WORK_DIR} \
|
|
--device ${DEVICE} \
|
|
--backend ${BACKEND}
|
|
```
|
|
|
|
|
|
<details>
|
|
<summary>Arguments</summary>
|
|
|
|
- `CONFIG_PATH`: the config file path of a detection method
|
|
|
|
- `WORK_DIR`: your path to save models and logs
|
|
|
|
- `MODEL_PATH`: the models to be quantized
|
|
|
|
- `DEVICE`: the device quantized models use (cpu/arm)
|
|
|
|
- `BACKEND`: the quantized models's framework (PyTorch/MNN)
|
|
|
|
</details>
|
|
|
|
**Examples:**
|
|
|
|
Edit `data_root`path in the `${CONFIG_PATH}` to your own data path.
|
|
|
|
```shell
|
|
python tools/quantize.py \
|
|
configs/detection/yolox/yolox_s_8xb16_300e_coco.py \
|
|
models/yolox_s.pth \
|
|
--device cpu \
|
|
--backend PyTorch
|
|
```
|
|
|
|
**Prune:**
|
|
This is used to prune yolox model; The pruned model will be saved in work_dir.
|
|
|
|
```shell
|
|
python tools/prune.py \
|
|
${CONFIG_PATH} \
|
|
${MODEL_PATH} \
|
|
--work_dir ${WORK_DIR} \
|
|
--pruning_class ${PRUNING_CLASS} \
|
|
--pruning_algorithm ${PRUNING_ALGORITHM}
|
|
```
|
|
|
|
|
|
<details>
|
|
<summary>Arguments</summary>
|
|
|
|
- `CONFIG_PATH`: the config file path of a detection method
|
|
|
|
- `WORK_DIR`: your path to save models and logs
|
|
|
|
- `MODEL_PATH`: the quantized models
|
|
|
|
- `PRUNING_CLASS`: pruning class for pruning models (AGP)
|
|
|
|
- `PRUNING_ALGORITHM`: pruning algorithm using by pruning class (taylorfo)
|
|
</details>
|
|
|
|
**Examples:**
|
|
|
|
Edit `data_root`path in the `${CONFIG_PATH}` to your own data path.
|
|
|
|
```shell
|
|
python tools/prune.py \
|
|
configs/detection/yolox/yolox_s_8xb16_300e_coco.py \
|
|
models/yolox_s.pth \
|
|
--pruning_class AGP \
|
|
--pruning_algorithm taylorfo
|
|
```
|
|
|
|
### Evaluation
|
|
|
|
Model will be auto-eval after compressing. So the result will be writed in the log.
|
|
|
|
### Inference
|
|
Download [test_image](http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/data/small_coco_demo/val2017/000000017627.jpg)
|
|
|
|
```python
|
|
import cv2
|
|
from easycv.predictors import TorchYoloXPredictor
|
|
|
|
output_ckpt = 'work_dirs/compression/yolox/quantize_model.pt'
|
|
detector = TorchYoloXPredictor(output_ckpt)
|
|
|
|
img = cv2.imread('000000017627.jpg')
|
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
output = detector.predict([img])
|
|
print(output)
|
|
|
|
# visualize image
|
|
from matplotlib import pyplot as plt
|
|
image = img.copy()
|
|
for box, cls_name in zip(output[0]['detection_boxes'], output[0]['detection_class_names']):
|
|
# box is [x1,y1,x2,y2]
|
|
box = [int(b) for b in box]
|
|
image = cv2.rectangle(image, tuple(box[:2]), tuple(box[2:4]), (0,255,0), 2)
|
|
cv2.putText(image, cls_name, (box[0], box[1]-5), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0,0,255), 2)
|
|
plt.imshow(image)
|
|
plt.show()
|
|
```
|