├── docs
├── .gitignore
├── demo
│ ├── ncnn_cpp_readme.md
│ ├── onnx_readme.md
│ ├── trt_cpp_readme.md
│ ├── trt_py_readme.md
│ ├── megengine_cpp_readme.md
│ ├── ncnn_android_readme.md
│ ├── openvino_cpp_readme.md
│ ├── openvino_py_readme.md
│ └── megengine_py_readme.md
├── requirements-doc.txt
├── index.rst
├── Makefile
├── _static
│ └── css
│ │ └── custom.css
├── freeze_module.md
├── assignment_visualization.md
├── updates_note.md
├── manipulate_training_image_size.md
├── mlflow_integration.md
├── quick_run.md
├── model_zoo.md
└── cache.md
├── demo
├── ncnn
│ ├── android
│ │ ├── settings.gradle
│ │ ├── gradle
│ │ │ └── wrapper
│ │ │ │ ├── gradle-wrapper.jar
│ │ │ │ └── gradle-wrapper.properties
│ │ ├── app
│ │ │ ├── src
│ │ │ │ └── main
│ │ │ │ │ ├── res
│ │ │ │ │ ├── values
│ │ │ │ │ │ └── strings.xml
│ │ │ │ │ └── layout
│ │ │ │ │ │ └── main.xml
│ │ │ │ │ ├── jni
│ │ │ │ │ └── CMakeLists.txt
│ │ │ │ │ ├── AndroidManifest.xml
│ │ │ │ │ └── java
│ │ │ │ │ └── com
│ │ │ │ │ └── megvii
│ │ │ │ │ └── yoloXncnn
│ │ │ │ │ ├── YOLOXncnn.java
│ │ │ │ │ └── yoloXncnn.java
│ │ │ └── build.gradle
│ │ ├── build.gradle
│ │ ├── README.md
│ │ └── gradlew.bat
│ ├── README.md
│ └── cpp
│ │ └── README.md
├── OpenVINO
│ ├── README.md
│ ├── cpp
│ │ ├── CMakeLists.txt
│ │ └── README.md
│ └── python
│ │ └── README.md
├── MegEngine
│ ├── python
│ │ ├── models
│ │ │ ├── __init__.py
│ │ │ ├── yolox.py
│ │ │ ├── yolo_fpn.py
│ │ │ └── yolo_pafpn.py
│ │ ├── README.md
│ │ ├── dump.py
│ │ ├── build.py
│ │ └── convert_weights.py
│ └── cpp
│ │ └── build.sh
├── TensorRT
│ ├── cpp
│ │ ├── CMakeLists.txt
│ │ └── README.md
│ └── python
│ │ └── README.md
├── nebullvm
│ ├── nebullvm_optimization.py
│ └── README.md
└── ONNXRuntime
│ ├── onnx_inference.py
│ └── README.md
├── tests
├── __init__.py
└── utils
│ └── test_model_utils.py
├── assets
├── demo.png
├── dog.jpg
├── logo.png
├── git_fig.png
├── sunjian.png
└── assignment.png
├── MANIFEST.in
├── yolox
├── __init__.py
├── core
│ ├── __init__.py
│ └── launch.py
├── exp
│ ├── __init__.py
│ ├── default
│ │ └── __init__.py
│ ├── build.py
│ └── base_exp.py
├── evaluators
│ └── __init__.py
├── models
│ ├── __init__.py
│ ├── yolox.py
│ ├── losses.py
│ ├── yolo_fpn.py
│ ├── yolo_pafpn.py
│ └── build.py
├── data
│ ├── datasets
│ │ ├── __init__.py
│ │ ├── voc_classes.py
│ │ └── coco_classes.py
│ ├── __init__.py
│ ├── data_prefetcher.py
│ ├── samplers.py
│ └── dataloading.py
├── utils
│ ├── compat.py
│ ├── __init__.py
│ ├── checkpoint.py
│ ├── ema.py
│ ├── setup_env.py
│ ├── allreduce_norm.py
│ ├── metric.py
│ ├── visualize.py
│ ├── boxes.py
│ └── demo_utils.py
├── layers
│ ├── __init__.py
│ ├── cocoeval
│ │ └── cocoeval.h
│ └── jit_ops.py
└── tools
│ └── __init__.py
├── tools
├── __init__.py
├── export_torchscript.py
├── trt.py
├── visualize_assign.py
├── export_onnx.py
└── train.py
├── exps
├── default
│ ├── __init__.py
│ ├── yolox_l.py
│ ├── yolox_m.py
│ ├── yolox_s.py
│ ├── yolox_x.py
│ ├── yolox_tiny.py
│ ├── yolov3.py
│ └── yolox_nano.py
└── example
│ ├── custom
│ ├── yolox_s.py
│ └── nano.py
│ └── yolox_voc
│ └── yolox_voc_s.py
├── requirements.txt
├── .github
└── workflows
│ ├── format_check.sh
│ └── ci.yaml
├── hubconf.py
├── .readthedocs.yaml
├── setup.cfg
├── datasets
└── README.md
├── .pre-commit-config.yaml
├── SECURITY.md
├── setup.py
└── .gitignore
/docs/.gitignore:
--------------------------------------------------------------------------------
1 | _build
--------------------------------------------------------------------------------
/demo/ncnn/android/settings.gradle:
--------------------------------------------------------------------------------
1 | include ':app'
2 |
--------------------------------------------------------------------------------
/docs/demo/ncnn_cpp_readme.md:
--------------------------------------------------------------------------------
1 | ../../demo/ncnn/cpp/README.md
--------------------------------------------------------------------------------
/docs/demo/onnx_readme.md:
--------------------------------------------------------------------------------
1 | ../../demo/ONNXRuntime/README.md
--------------------------------------------------------------------------------
/docs/demo/trt_cpp_readme.md:
--------------------------------------------------------------------------------
1 | ../../demo/TensorRT/cpp/README.md
--------------------------------------------------------------------------------
/docs/demo/trt_py_readme.md:
--------------------------------------------------------------------------------
1 | ../../demo/TensorRT/python/README.md
--------------------------------------------------------------------------------
/docs/demo/megengine_cpp_readme.md:
--------------------------------------------------------------------------------
1 | ../../demo/MegEngine/cpp/README.md
--------------------------------------------------------------------------------
/docs/demo/ncnn_android_readme.md:
--------------------------------------------------------------------------------
1 | ../../demo/ncnn/android/README.md
--------------------------------------------------------------------------------
/docs/demo/openvino_cpp_readme.md:
--------------------------------------------------------------------------------
1 | ../../demo/OpenVINO/cpp/README.md
--------------------------------------------------------------------------------
/docs/demo/openvino_py_readme.md:
--------------------------------------------------------------------------------
1 | ../../demo/OpenVINO/python/README.md
--------------------------------------------------------------------------------
/docs/demo/megengine_py_readme.md:
--------------------------------------------------------------------------------
1 | ../../demo/MegEngine/python/README.md
--------------------------------------------------------------------------------
/tests/__init__.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 |
--------------------------------------------------------------------------------
/assets/demo.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Megvii-BaseDetection/YOLOX/HEAD/assets/demo.png
--------------------------------------------------------------------------------
/assets/dog.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Megvii-BaseDetection/YOLOX/HEAD/assets/dog.jpg
--------------------------------------------------------------------------------
/assets/logo.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Megvii-BaseDetection/YOLOX/HEAD/assets/logo.png
--------------------------------------------------------------------------------
/MANIFEST.in:
--------------------------------------------------------------------------------
1 | include requirements.txt
2 | recursive-include yolox *.cpp *.h *.cu *.cuh *.cc
3 |
--------------------------------------------------------------------------------
/assets/git_fig.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Megvii-BaseDetection/YOLOX/HEAD/assets/git_fig.png
--------------------------------------------------------------------------------
/assets/sunjian.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Megvii-BaseDetection/YOLOX/HEAD/assets/sunjian.png
--------------------------------------------------------------------------------
/demo/OpenVINO/README.md:
--------------------------------------------------------------------------------
1 | ## YOLOX for OpenVINO
2 |
3 | * [C++ Demo](./cpp)
4 | * [Python Demo](./python)
--------------------------------------------------------------------------------
/yolox/__init__.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 |
4 | __version__ = "0.3.0"
5 |
--------------------------------------------------------------------------------
/assets/assignment.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Megvii-BaseDetection/YOLOX/HEAD/assets/assignment.png
--------------------------------------------------------------------------------
/tools/__init__.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
--------------------------------------------------------------------------------
/exps/default/__init__.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
--------------------------------------------------------------------------------
/demo/ncnn/android/gradle/wrapper/gradle-wrapper.jar:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Megvii-BaseDetection/YOLOX/HEAD/demo/ncnn/android/gradle/wrapper/gradle-wrapper.jar
--------------------------------------------------------------------------------
/demo/ncnn/android/app/src/main/res/values/strings.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 | yoloXncnn
4 |
5 |
--------------------------------------------------------------------------------
/yolox/core/__init__.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
5 | from .launch import launch
6 | from .trainer import Trainer
7 |
--------------------------------------------------------------------------------
/yolox/exp/__init__.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # Copyright (c) Megvii Inc. All rights reserved.
3 |
4 | from .base_exp import BaseExp
5 | from .build import get_exp
6 | from .yolox_base import Exp, check_exp_value
7 |
--------------------------------------------------------------------------------
/yolox/evaluators/__init__.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
5 | from .coco_evaluator import COCOEvaluator
6 | from .voc_evaluator import VOCEvaluator
7 |
--------------------------------------------------------------------------------
/docs/requirements-doc.txt:
--------------------------------------------------------------------------------
1 | docutils==0.16
2 | # https://github.com/sphinx-doc/sphinx/commit/7acd3ada3f38076af7b2b5c9f3b60bb9c2587a3d
3 | sphinx==3.2.0
4 | recommonmark==0.6.0
5 | sphinx_rtd_theme
6 | omegaconf>=2.1.0.dev24
7 | hydra-core>=1.1.0.dev5
8 | sphinx-markdown-tables==0.0.15
9 |
--------------------------------------------------------------------------------
/demo/ncnn/README.md:
--------------------------------------------------------------------------------
1 | # YOLOX-ncnn
2 |
3 | Compile files of YOLOX object detection base on [ncnn](https://github.com/Tencent/ncnn).
4 | YOLOX is included in ncnn now, you could also try building from ncnn, it's better.
5 |
6 | ## Acknowledgement
7 |
8 | * [ncnn](https://github.com/Tencent/ncnn)
9 |
--------------------------------------------------------------------------------
/demo/ncnn/android/gradle/wrapper/gradle-wrapper.properties:
--------------------------------------------------------------------------------
1 | #Sun Aug 25 10:34:48 CST 2019
2 | distributionBase=GRADLE_USER_HOME
3 | distributionPath=wrapper/dists
4 | zipStoreBase=GRADLE_USER_HOME
5 | zipStorePath=wrapper/dists
6 | distributionUrl=https\://services.gradle.org/distributions/gradle-5.4.1-all.zip
7 |
--------------------------------------------------------------------------------
/demo/MegEngine/python/models/__init__.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii Inc. All rights reserved.
4 |
5 | from .darknet import CSPDarknet, Darknet
6 | from .yolo_fpn import YOLOFPN
7 | from .yolo_head import YOLOXHead
8 | from .yolo_pafpn import YOLOPAFPN
9 | from .yolox import YOLOX
10 |
--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | # TODO: Update with exact module version
2 | numpy
3 | torch>=1.7
4 | opencv_python
5 | loguru
6 | tqdm
7 | torchvision
8 | thop
9 | ninja
10 | tabulate
11 | psutil
12 | tensorboard
13 |
14 | # verified versions
15 | # pycocotools corresponds to https://github.com/ppwwyyxx/cocoapi
16 | pycocotools>=2.0.2
17 | onnx>=1.13.0
18 | onnx-simplifier==0.4.10
19 |
--------------------------------------------------------------------------------
/yolox/models/__init__.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii Inc. All rights reserved.
4 |
5 | from .build import *
6 | from .darknet import CSPDarknet, Darknet
7 | from .losses import IOUloss
8 | from .yolo_fpn import YOLOFPN
9 | from .yolo_head import YOLOXHead
10 | from .yolo_pafpn import YOLOPAFPN
11 | from .yolox import YOLOX
12 |
--------------------------------------------------------------------------------
/yolox/data/datasets/__init__.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
5 | from .coco import COCODataset
6 | from .coco_classes import COCO_CLASSES
7 | from .datasets_wrapper import CacheDataset, ConcatDataset, Dataset, MixConcatDataset
8 | from .mosaicdetection import MosaicDetection
9 | from .voc import VOCDetection
10 |
--------------------------------------------------------------------------------
/yolox/utils/compat.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 |
4 | import torch
5 |
6 | _TORCH_VER = [int(x) for x in torch.__version__.split(".")[:2]]
7 |
8 | __all__ = ["meshgrid"]
9 |
10 |
11 | def meshgrid(*tensors):
12 | if _TORCH_VER >= [1, 10]:
13 | return torch.meshgrid(*tensors, indexing="ij")
14 | else:
15 | return torch.meshgrid(*tensors)
16 |
--------------------------------------------------------------------------------
/demo/ncnn/android/app/src/main/jni/CMakeLists.txt:
--------------------------------------------------------------------------------
1 | project(yoloXncnn)
2 |
3 | cmake_minimum_required(VERSION 3.4.1)
4 |
5 | set(ncnn_DIR ${CMAKE_SOURCE_DIR}/ncnn-20210525-android-vulkan/${ANDROID_ABI}/lib/cmake/ncnn)
6 | find_package(ncnn REQUIRED)
7 |
8 | add_library(yoloXncnn SHARED yoloXncnn_jni.cpp)
9 |
10 | target_link_libraries(yoloXncnn
11 | ncnn
12 |
13 | jnigraphics
14 | )
15 |
--------------------------------------------------------------------------------
/yolox/data/__init__.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
5 | from .data_augment import TrainTransform, ValTransform
6 | from .data_prefetcher import DataPrefetcher
7 | from .dataloading import DataLoader, get_yolox_datadir, worker_init_reset_seed
8 | from .datasets import *
9 | from .samplers import InfiniteSampler, YoloBatchSampler
10 |
--------------------------------------------------------------------------------
/.github/workflows/format_check.sh:
--------------------------------------------------------------------------------
1 | #!/bin/bash -e
2 |
3 | set -e
4 |
5 | export PYTHONPATH=$PWD:$PYTHONPATH
6 |
7 | flake8 yolox exps tools || flake8_ret=$?
8 | if [ "$flake8_ret" ]; then
9 | exit $flake8_ret
10 | fi
11 | echo "All flake check passed!"
12 | isort --check-only -rc yolox exps || isort_ret=$?
13 | if [ "$isort_ret" ]; then
14 | exit $isort_ret
15 | fi
16 | echo "All isort check passed!"
17 |
--------------------------------------------------------------------------------
/demo/ncnn/android/build.gradle:
--------------------------------------------------------------------------------
1 | // Top-level build file where you can add configuration options common to all sub-projects/modules.
2 | buildscript {
3 | repositories {
4 | jcenter()
5 | google()
6 | }
7 | dependencies {
8 | classpath 'com.android.tools.build:gradle:3.5.0'
9 | }
10 | }
11 |
12 | allprojects {
13 | repositories {
14 | jcenter()
15 | google()
16 | }
17 | }
18 |
--------------------------------------------------------------------------------
/exps/default/yolox_l.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
5 | import os
6 |
7 | from yolox.exp import Exp as MyExp
8 |
9 |
10 | class Exp(MyExp):
11 | def __init__(self):
12 | super(Exp, self).__init__()
13 | self.depth = 1.0
14 | self.width = 1.0
15 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
16 |
--------------------------------------------------------------------------------
/exps/default/yolox_m.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
5 | import os
6 |
7 | from yolox.exp import Exp as MyExp
8 |
9 |
10 | class Exp(MyExp):
11 | def __init__(self):
12 | super(Exp, self).__init__()
13 | self.depth = 0.67
14 | self.width = 0.75
15 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
16 |
--------------------------------------------------------------------------------
/exps/default/yolox_s.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
5 | import os
6 |
7 | from yolox.exp import Exp as MyExp
8 |
9 |
10 | class Exp(MyExp):
11 | def __init__(self):
12 | super(Exp, self).__init__()
13 | self.depth = 0.33
14 | self.width = 0.50
15 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
16 |
--------------------------------------------------------------------------------
/exps/default/yolox_x.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
5 | import os
6 |
7 | from yolox.exp import Exp as MyExp
8 |
9 |
10 | class Exp(MyExp):
11 | def __init__(self):
12 | super(Exp, self).__init__()
13 | self.depth = 1.33
14 | self.width = 1.25
15 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
16 |
--------------------------------------------------------------------------------
/yolox/layers/__init__.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii Inc. All rights reserved.
4 |
5 | # import torch first to make jit op work without `ImportError of libc10.so`
6 | import torch # noqa
7 |
8 | from .jit_ops import FastCOCOEvalOp, JitOp
9 |
10 | try:
11 | from .fast_coco_eval_api import COCOeval_opt
12 | except ImportError: # exception will be raised when users build yolox from source
13 | pass
14 |
--------------------------------------------------------------------------------
/yolox/utils/__init__.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # Copyright (c) Megvii Inc. All rights reserved.
3 |
4 | from .allreduce_norm import *
5 | from .boxes import *
6 | from .checkpoint import load_ckpt, save_checkpoint
7 | from .compat import meshgrid
8 | from .demo_utils import *
9 | from .dist import *
10 | from .ema import *
11 | from .logger import WandbLogger, setup_logger
12 | from .lr_scheduler import LRScheduler
13 | from .metric import *
14 | from .mlflow_logger import MlflowLogger
15 | from .model_utils import *
16 | from .setup_env import *
17 | from .visualize import *
18 |
--------------------------------------------------------------------------------
/yolox/data/datasets/voc_classes.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
5 | # VOC_CLASSES = ( '__background__', # always index 0
6 | VOC_CLASSES = (
7 | "aeroplane",
8 | "bicycle",
9 | "bird",
10 | "boat",
11 | "bottle",
12 | "bus",
13 | "car",
14 | "cat",
15 | "chair",
16 | "cow",
17 | "diningtable",
18 | "dog",
19 | "horse",
20 | "motorbike",
21 | "person",
22 | "pottedplant",
23 | "sheep",
24 | "sofa",
25 | "train",
26 | "tvmonitor",
27 | )
28 |
--------------------------------------------------------------------------------
/hubconf.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 |
4 | """
5 | Usage example:
6 | import torch
7 | model = torch.hub.load("Megvii-BaseDetection/YOLOX", "yolox_s")
8 | model = torch.hub.load("Megvii-BaseDetection/YOLOX", "yolox_custom",
9 | exp_path="exp.py", ckpt_path="ckpt.pth")
10 | """
11 | dependencies = ["torch"]
12 |
13 | from yolox.models import ( # isort:skip # noqa: F401, E402
14 | yolox_tiny,
15 | yolox_nano,
16 | yolox_s,
17 | yolox_m,
18 | yolox_l,
19 | yolox_x,
20 | yolov3,
21 | yolox_custom
22 | )
23 |
--------------------------------------------------------------------------------
/demo/OpenVINO/cpp/CMakeLists.txt:
--------------------------------------------------------------------------------
1 | cmake_minimum_required(VERSION 3.4.1)
2 | set(CMAKE_CXX_STANDARD 14)
3 |
4 | project(yolox_openvino_demo)
5 |
6 | find_package(OpenCV REQUIRED)
7 | find_package(InferenceEngine REQUIRED)
8 | find_package(ngraph REQUIRED)
9 |
10 | include_directories(
11 | ${OpenCV_INCLUDE_DIRS}
12 | ${CMAKE_CURRENT_SOURCE_DIR}
13 | ${CMAKE_CURRENT_BINARY_DIR}
14 | )
15 |
16 | add_executable(yolox_openvino yolox_openvino.cpp)
17 |
18 | target_link_libraries(
19 | yolox_openvino
20 | ${InferenceEngine_LIBRARIES}
21 | ${NGRAPH_LIBRARIES}
22 | ${OpenCV_LIBS}
23 | )
--------------------------------------------------------------------------------
/demo/ncnn/android/app/build.gradle:
--------------------------------------------------------------------------------
1 | apply plugin: 'com.android.application'
2 |
3 | android {
4 | compileSdkVersion 24
5 | buildToolsVersion "29.0.2"
6 |
7 | defaultConfig {
8 | applicationId "com.megvii.yoloXncnn"
9 | archivesBaseName = "$applicationId"
10 |
11 | ndk {
12 | moduleName "ncnn"
13 | abiFilters "armeabi-v7a", "arm64-v8a"
14 | }
15 | minSdkVersion 24
16 | }
17 |
18 | externalNativeBuild {
19 | cmake {
20 | version "3.10.2"
21 | path file('src/main/jni/CMakeLists.txt')
22 | }
23 | }
24 | }
25 |
--------------------------------------------------------------------------------
/.readthedocs.yaml:
--------------------------------------------------------------------------------
1 | # .readthedocs.yaml
2 | # Read the Docs configuration file
3 | # See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
4 |
5 | # Required
6 | version: 2
7 |
8 | # Build documentation in the docs/ directory with Sphinx
9 | sphinx:
10 | configuration: docs/conf.py
11 |
12 | # Optionally build your docs in additional formats such as PDF
13 | formats:
14 | - pdf
15 |
16 | # Optionally set the version of Python and requirements required to build your docs
17 | python:
18 | version: "3.7"
19 | install:
20 | - requirements: docs/requirements-doc.txt
21 | - requirements: requirements.txt
22 |
--------------------------------------------------------------------------------
/exps/default/yolox_tiny.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
5 | import os
6 |
7 | from yolox.exp import Exp as MyExp
8 |
9 |
10 | class Exp(MyExp):
11 | def __init__(self):
12 | super(Exp, self).__init__()
13 | self.depth = 0.33
14 | self.width = 0.375
15 | self.input_size = (416, 416)
16 | self.mosaic_scale = (0.5, 1.5)
17 | self.random_size = (10, 20)
18 | self.test_size = (416, 416)
19 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
20 | self.enable_mixup = False
21 |
--------------------------------------------------------------------------------
/setup.cfg:
--------------------------------------------------------------------------------
1 | [isort]
2 | line_length = 100
3 | multi_line_output = 3
4 | balanced_wrapping = True
5 | known_standard_library = setuptools
6 | known_third_party = tqdm,loguru,tabulate,psutil
7 | known_data_processing = cv2,numpy,scipy,PIL,matplotlib
8 | known_datasets = pycocotools
9 | known_deeplearning = torch,torchvision,caffe2,onnx,apex,timm,thop,torch2trt,tensorrt,openvino,onnxruntime
10 | known_myself = yolox
11 | sections = FUTURE,STDLIB,THIRDPARTY,data_processing,datasets,deeplearning,myself,FIRSTPARTY,LOCALFOLDER
12 | no_lines_before=STDLIB,THIRDPARTY,datasets
13 | default_section = FIRSTPARTY
14 |
15 | [flake8]
16 | max-line-length = 100
17 | max-complexity = 18
18 | exclude = __init__.py
19 |
--------------------------------------------------------------------------------
/docs/index.rst:
--------------------------------------------------------------------------------
1 |
2 | Welcome to YOLOX's documentation!
3 | ======================================
4 |
5 | .. image:: ../assets/logo.png
6 |
7 | .. toctree::
8 | :maxdepth: 2
9 | :caption: Quick Run
10 |
11 | quick_run
12 | model_zoo
13 |
14 | .. toctree::
15 | :maxdepth: 2
16 | :caption: Tutorials
17 |
18 | train_custom_data
19 |
20 | .. toctree::
21 | :maxdepth: 2
22 | :caption: Deployment
23 |
24 | demo/trt_py_readme
25 | demo/trt_cpp_readme
26 | demo/megengine_cpp_readme
27 | demo/megengine_py_readme
28 | demo/ncnn_android_readme
29 | demo/ncnn_cpp_readme
30 | demo/onnx_readme
31 | demo/openvino_py_readme
32 | demo/openvino_cpp_readme
33 |
--------------------------------------------------------------------------------
/docs/Makefile:
--------------------------------------------------------------------------------
1 | # Minimal makefile for Sphinx documentation
2 | # Copyright (c) Facebook, Inc. and its affiliates.
3 |
4 | # You can set these variables from the command line.
5 | SPHINXOPTS =
6 | SPHINXBUILD = sphinx-build
7 | SOURCEDIR = .
8 | BUILDDIR = _build
9 |
10 | # Put it first so that "make" without argument is like "make help".
11 | help:
12 | @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
13 |
14 | .PHONY: help Makefile
15 |
16 | # Catch-all target: route all unknown targets to Sphinx using the new
17 | # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS).
18 | %: Makefile
19 | @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
--------------------------------------------------------------------------------
/docs/_static/css/custom.css:
--------------------------------------------------------------------------------
1 | /*
2 | * Copyright (c) Facebook, Inc. and its affiliates.
3 | * some extra css to make markdown look similar between github/sphinx
4 | */
5 |
6 | /*
7 | * Below is for install.md:
8 | */
9 | .rst-content code {
10 | white-space: pre;
11 | border: 0px;
12 | }
13 |
14 | .rst-content th {
15 | border: 1px solid #e1e4e5;
16 | }
17 |
18 | .rst-content th p {
19 | /* otherwise will be default 24px for regular paragraph */
20 | margin-bottom: 0px;
21 | }
22 |
23 | .rst-content .line-block {
24 | /* otherwise will be 24px */
25 | margin-bottom: 0px;
26 | }
27 |
28 | div.section > details {
29 | padding-bottom: 1em;
30 | }
31 |
--------------------------------------------------------------------------------
/demo/ncnn/android/app/src/main/AndroidManifest.xml:
--------------------------------------------------------------------------------
1 |
2 |
6 |
7 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
--------------------------------------------------------------------------------
/demo/ncnn/android/app/src/main/java/com/megvii/yoloXncnn/YOLOXncnn.java:
--------------------------------------------------------------------------------
1 | // Copyright (C) Megvii, Inc. and its affiliates. All rights reserved.
2 |
3 | package com.megvii.yoloXncnn;
4 |
5 | import android.content.res.AssetManager;
6 | import android.graphics.Bitmap;
7 |
8 | public class YOLOXncnn
9 | {
10 | public native boolean Init(AssetManager mgr);
11 |
12 | public class Obj
13 | {
14 | public float x;
15 | public float y;
16 | public float w;
17 | public float h;
18 | public String label;
19 | public float prob;
20 | }
21 |
22 | public native Obj[] Detect(Bitmap bitmap, boolean use_gpu);
23 |
24 | static {
25 | System.loadLibrary("yoloXncnn");
26 | }
27 | }
28 |
--------------------------------------------------------------------------------
/demo/ncnn/android/app/src/main/java/com/megvii/yoloXncnn/yoloXncnn.java:
--------------------------------------------------------------------------------
1 | // Copyright (C) Megvii, Inc. and its affiliates. All rights reserved.
2 |
3 | package com.megvii.yoloXncnn;
4 |
5 | import android.content.res.AssetManager;
6 | import android.graphics.Bitmap;
7 |
8 | public class YOLOXncnn
9 | {
10 | public native boolean Init(AssetManager mgr);
11 |
12 | public class Obj
13 | {
14 | public float x;
15 | public float y;
16 | public float w;
17 | public float h;
18 | public String label;
19 | public float prob;
20 | }
21 |
22 | public native Obj[] Detect(Bitmap bitmap, boolean use_gpu);
23 |
24 | static {
25 | System.loadLibrary("yoloXncnn");
26 | }
27 | }
28 |
--------------------------------------------------------------------------------
/exps/example/custom/yolox_s.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 | import os
5 |
6 | from yolox.exp import Exp as MyExp
7 |
8 |
9 | class Exp(MyExp):
10 | def __init__(self):
11 | super(Exp, self).__init__()
12 | self.depth = 0.33
13 | self.width = 0.50
14 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
15 |
16 | # Define yourself dataset path
17 | self.data_dir = "datasets/coco128"
18 | self.train_ann = "instances_train2017.json"
19 | self.val_ann = "instances_val2017.json"
20 |
21 | self.num_classes = 71
22 |
23 | self.max_epoch = 300
24 | self.data_num_workers = 4
25 | self.eval_interval = 1
26 |
--------------------------------------------------------------------------------
/datasets/README.md:
--------------------------------------------------------------------------------
1 | # Prepare datasets
2 |
3 | If you have a dataset directory, you could use os environment variable named `YOLOX_DATADIR`. Under this directory, YOLOX will look for datasets in the structure described below, if needed.
4 | ```
5 | $YOLOX_DATADIR/
6 | COCO/
7 | ```
8 | You can set the location for builtin datasets by
9 | ```shell
10 | export YOLOX_DATADIR=/path/to/your/datasets
11 | ```
12 | If `YOLOX_DATADIR` is not set, the default value of dataset directory is `./datasets` relative to your current working directory.
13 |
14 | ## Expected dataset structure for [COCO detection](https://cocodataset.org/#download):
15 |
16 | ```
17 | COCO/
18 | annotations/
19 | instances_{train,val}2017.json
20 | {train,val}2017/
21 | # image files that are mentioned in the corresponding json
22 | ```
23 |
24 | You can use the 2014 version of the dataset as well.
25 |
--------------------------------------------------------------------------------
/demo/MegEngine/python/README.md:
--------------------------------------------------------------------------------
1 | # YOLOX-Python-MegEngine
2 |
3 | Python version of YOLOX object detection base on [MegEngine](https://github.com/MegEngine/MegEngine).
4 |
5 | ## Tutorial
6 |
7 | ### Step1: install requirements
8 |
9 | ```
10 | python3 -m pip install megengine -f https://megengine.org.cn/whl/mge.html
11 | ```
12 |
13 | ### Step2: convert checkpoint weights from torch's path file
14 |
15 | ```
16 | python3 convert_weights.py -w yolox_s.pth -o yolox_s_mge.pkl
17 | ```
18 |
19 | ### Step3: run demo
20 |
21 | This part is the same as torch's python demo, but no need to specify device.
22 |
23 | ```
24 | python3 demo.py image -n yolox-s -c yolox_s_mge.pkl --path ../../../assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result
25 | ```
26 |
27 | ### [Optional]Step4: dump model for cpp inference
28 |
29 | > **Note**: result model is dumped with `optimize_for_inference` and `enable_fuse_conv_bias_nonlinearity`.
30 |
31 | ```
32 | python3 dump.py -n yolox-s -c yolox_s_mge.pkl --dump_path yolox_s.mge
33 | ```
34 |
--------------------------------------------------------------------------------
/demo/MegEngine/python/models/yolox.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- encoding: utf-8 -*-
3 | # Copyright (c) Megvii Inc. All rights reserved.
4 |
5 | import megengine.module as M
6 |
7 | from .yolo_head import YOLOXHead
8 | from .yolo_pafpn import YOLOPAFPN
9 |
10 |
11 | class YOLOX(M.Module):
12 | """
13 | YOLOX model module. The module list is defined by create_yolov3_modules function.
14 | The network returns loss values from three YOLO layers during training
15 | and detection results during test.
16 | """
17 |
18 | def __init__(self, backbone=None, head=None):
19 | super().__init__()
20 | if backbone is None:
21 | backbone = YOLOPAFPN()
22 | if head is None:
23 | head = YOLOXHead(80)
24 |
25 | self.backbone = backbone
26 | self.head = head
27 |
28 | def forward(self, x):
29 | # fpn output content features of [dark3, dark4, dark5]
30 | fpn_outs = self.backbone(x)
31 | assert not self.training
32 | outputs = self.head(fpn_outs)
33 |
34 | return outputs
35 |
--------------------------------------------------------------------------------
/yolox/tools/__init__.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # Copyright (c) Megvii Inc. All rights reserved.
3 |
4 | # This file is used for package installation. Script of train/eval/export will be available.
5 |
6 | import sys
7 | from importlib import abc, util
8 | from pathlib import Path
9 |
10 | _TOOLS_PATH = Path(__file__).resolve().parent.parent.parent / "tools"
11 |
12 | if _TOOLS_PATH.is_dir():
13 | # This is true only for in-place installation (pip install -e, setup.py develop),
14 | # where setup(package_dir=) does not work: https://github.com/pypa/setuptools/issues/230
15 |
16 | class _PathFinder(abc.MetaPathFinder):
17 |
18 | def find_spec(self, name, path, target=None):
19 | if not name.startswith("yolox.tools."):
20 | return
21 | project_name = name.split(".")[-1] + ".py"
22 | target_file = _TOOLS_PATH / project_name
23 | if not target_file.is_file():
24 | return
25 | return util.spec_from_file_location(name, target_file)
26 |
27 | sys.meta_path.append(_PathFinder())
28 |
--------------------------------------------------------------------------------
/yolox/exp/default/__init__.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii Inc. All rights reserved.
4 |
5 | # This file is used for package installation and find default exp file
6 |
7 | import sys
8 | from importlib import abc, util
9 | from pathlib import Path
10 |
11 | _EXP_PATH = Path(__file__).resolve().parent.parent.parent.parent / "exps" / "default"
12 |
13 | if _EXP_PATH.is_dir():
14 | # This is true only for in-place installation (pip install -e, setup.py develop),
15 | # where setup(package_dir=) does not work: https://github.com/pypa/setuptools/issues/230
16 |
17 | class _ExpFinder(abc.MetaPathFinder):
18 |
19 | def find_spec(self, name, path, target=None):
20 | if not name.startswith("yolox.exp.default"):
21 | return
22 | project_name = name.split(".")[-1] + ".py"
23 | target_file = _EXP_PATH / project_name
24 | if not target_file.is_file():
25 | return
26 | return util.spec_from_file_location(name, target_file)
27 |
28 | sys.meta_path.append(_ExpFinder())
29 |
--------------------------------------------------------------------------------
/exps/default/yolov3.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
5 | import os
6 |
7 | import torch.nn as nn
8 |
9 | from yolox.exp import Exp as MyExp
10 |
11 |
12 | class Exp(MyExp):
13 | def __init__(self):
14 | super(Exp, self).__init__()
15 | self.depth = 1.0
16 | self.width = 1.0
17 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
18 |
19 | def get_model(self, sublinear=False):
20 | def init_yolo(M):
21 | for m in M.modules():
22 | if isinstance(m, nn.BatchNorm2d):
23 | m.eps = 1e-3
24 | m.momentum = 0.03
25 | if "model" not in self.__dict__:
26 | from yolox.models import YOLOX, YOLOFPN, YOLOXHead
27 | backbone = YOLOFPN()
28 | head = YOLOXHead(self.num_classes, self.width, in_channels=[128, 256, 512], act="lrelu")
29 | self.model = YOLOX(backbone, head)
30 | self.model.apply(init_yolo)
31 | self.model.head.initialize_biases(1e-2)
32 |
33 | return self.model
34 |
--------------------------------------------------------------------------------
/demo/ncnn/android/README.md:
--------------------------------------------------------------------------------
1 | # YOLOX-Android-ncnn
2 |
3 | Andoird app of YOLOX object detection base on [ncnn](https://github.com/Tencent/ncnn)
4 |
5 |
6 | ## Tutorial
7 |
8 | ### Step1
9 |
10 | Download ncnn-android-vulkan.zip from [releases of ncnn](https://github.com/Tencent/ncnn/releases). This repo uses
11 | [20210525 release](https://github.com/Tencent/ncnn/releases/download/20210525/ncnn-20210525-android-vulkan.zip) for building.
12 |
13 | ### Step2
14 |
15 | After downloading, please extract your zip file. Then, there are two ways to finish this step:
16 | * put your extracted directory into **app/src/main/jni**
17 | * change the **ncnn_DIR** path in **app/src/main/jni/CMakeLists.txt** to your extracted directory
18 |
19 | ### Step3
20 | Download example param and bin file from [onedrive](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/ESXBH_GSSmFMszWJ6YG2VkQB5cWDfqVWXgk0D996jH0rpQ?e=qzEqUh) or [github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_s_ncnn.tar.gz). Unzip the file to **app/src/main/assets**.
21 |
22 | ### Step4
23 | Open this project with Android Studio, build it and enjoy!
24 |
25 | ## Reference
26 |
27 | * [ncnn-android-yolov5](https://github.com/nihui/ncnn-android-yolov5)
28 |
--------------------------------------------------------------------------------
/docs/freeze_module.md:
--------------------------------------------------------------------------------
1 | # Freeze module
2 |
3 | This page guide users to freeze module in YOLOX.
4 | Exp controls everything in YOLOX, so let's start from creating an Exp object.
5 |
6 | ## 1. Create your own expermiment object
7 |
8 | We take an example of YOLOX-S model on COCO dataset to give a more clear guide.
9 |
10 | Import the config you want (or write your own Exp object inherit from `yolox.exp.BaseExp`).
11 | ```python
12 | from yolox.exp.default.yolox_s import Exp as MyExp
13 | ```
14 |
15 | ## 2. Override `get_model` method
16 |
17 | Here is a simple code to freeze backbone (FPN not included) of module.
18 | ```python
19 | class Exp(MyExp):
20 |
21 | def get_model(self):
22 | from yolox.utils import freeze_module
23 | model = super().get_model()
24 | freeze_module(model.backbone.backbone)
25 | return model
26 | ```
27 | if you only want to freeze FPN, `freeze_module(model.backbone)` might help.
28 |
29 | ## 3. Train
30 | Suppose that the path of your Exp is `/path/to/my_exp.py`, use the following command to train your model.
31 | ```bash
32 | python3 -m yolox.tools.train -f /path/to/my_exp.py
33 | ```
34 | For more details of training, run the following command.
35 | ```bash
36 | python3 -m yolox.tools.train --help
37 | ```
38 |
--------------------------------------------------------------------------------
/demo/TensorRT/cpp/CMakeLists.txt:
--------------------------------------------------------------------------------
1 | cmake_minimum_required(VERSION 2.6)
2 |
3 | project(yolox)
4 |
5 | add_definitions(-std=c++11)
6 |
7 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
8 | set(CMAKE_CXX_STANDARD 11)
9 | set(CMAKE_BUILD_TYPE Debug)
10 |
11 | find_package(CUDA REQUIRED)
12 |
13 | include_directories(${PROJECT_SOURCE_DIR}/include)
14 | # include and link dirs of cuda and tensorrt, you need adapt them if yours are different
15 | # cuda
16 | include_directories(/data/cuda/cuda-10.2/cuda/include)
17 | link_directories(/data/cuda/cuda-10.2/cuda/lib64)
18 | # cudnn
19 | include_directories(/data/cuda/cuda-10.2/cudnn/v8.0.4/include)
20 | link_directories(/data/cuda/cuda-10.2/cudnn/v8.0.4/lib64)
21 | # tensorrt
22 | include_directories(/data/cuda/cuda-10.2/TensorRT/v7.2.1.6/include)
23 | link_directories(/data/cuda/cuda-10.2/TensorRT/v7.2.1.6/lib)
24 |
25 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED")
26 |
27 | find_package(OpenCV)
28 | include_directories(${OpenCV_INCLUDE_DIRS})
29 |
30 | add_executable(yolox ${PROJECT_SOURCE_DIR}/yolox.cpp)
31 | target_link_libraries(yolox nvinfer)
32 | target_link_libraries(yolox cudart)
33 | target_link_libraries(yolox ${OpenCV_LIBS})
34 |
35 | add_definitions(-O2 -pthread)
36 |
37 |
--------------------------------------------------------------------------------
/demo/ncnn/android/app/src/main/res/layout/main.xml:
--------------------------------------------------------------------------------
1 |
2 |
6 |
7 |
11 |
12 |
17 |
22 |
27 |
28 |
29 |
34 |
35 |
36 |
--------------------------------------------------------------------------------
/.pre-commit-config.yaml:
--------------------------------------------------------------------------------
1 | repos:
2 | - repo: https://github.com/pycqa/flake8
3 | rev: 3.8.3
4 | hooks:
5 | - id: flake8
6 | - repo: https://github.com/pre-commit/pre-commit-hooks
7 | rev: v3.1.0
8 | hooks:
9 | - id: check-added-large-files
10 | - id: check-docstring-first
11 | - id: check-executables-have-shebangs
12 | - id: check-json
13 | - id: check-yaml
14 | args: ["--unsafe"]
15 | - id: debug-statements
16 | - id: end-of-file-fixer
17 | - id: requirements-txt-fixer
18 | - id: trailing-whitespace
19 | - repo: https://github.com/jorisroovers/gitlint
20 | rev: v0.15.1
21 | hooks:
22 | - id: gitlint
23 | - repo: https://github.com/pycqa/isort
24 | rev: 4.3.21
25 | hooks:
26 | - id: isort
27 |
28 | - repo: https://github.com/PyCQA/autoflake
29 | rev: v1.4
30 | hooks:
31 | - id: autoflake
32 | name: Remove unused variables and imports
33 | entry: autoflake
34 | language: python
35 | args:
36 | [
37 | "--in-place",
38 | "--remove-all-unused-imports",
39 | "--remove-unused-variables",
40 | "--expand-star-imports",
41 | "--ignore-init-module-imports",
42 | ]
43 | files: \.py$
44 |
--------------------------------------------------------------------------------
/yolox/exp/build.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii Inc. All rights reserved.
4 |
5 | import importlib
6 | import os
7 | import sys
8 |
9 |
10 | def get_exp_by_file(exp_file):
11 | try:
12 | sys.path.append(os.path.dirname(exp_file))
13 | current_exp = importlib.import_module(os.path.basename(exp_file).split(".")[0])
14 | exp = current_exp.Exp()
15 | except Exception:
16 | raise ImportError("{} doesn't contains class named 'Exp'".format(exp_file))
17 | return exp
18 |
19 |
20 | def get_exp_by_name(exp_name):
21 | exp = exp_name.replace("-", "_") # convert string like "yolox-s" to "yolox_s"
22 | module_name = ".".join(["yolox", "exp", "default", exp])
23 | exp_object = importlib.import_module(module_name).Exp()
24 | return exp_object
25 |
26 |
27 | def get_exp(exp_file=None, exp_name=None):
28 | """
29 | get Exp object by file or name. If exp_file and exp_name
30 | are both provided, get Exp by exp_file.
31 |
32 | Args:
33 | exp_file (str): file path of experiment.
34 | exp_name (str): name of experiment. "yolo-s",
35 | """
36 | assert (
37 | exp_file is not None or exp_name is not None
38 | ), "plz provide exp file or exp name."
39 | if exp_file is not None:
40 | return get_exp_by_file(exp_file)
41 | else:
42 | return get_exp_by_name(exp_name)
43 |
--------------------------------------------------------------------------------
/.github/workflows/ci.yaml:
--------------------------------------------------------------------------------
1 | # This is a basic workflow to help you get started with Actions
2 |
3 | name: CI
4 |
5 | # Controls when the action will run. Triggers the workflow on push or pull request
6 | # events but only for the master branch
7 | on:
8 | push:
9 | pull_request:
10 |
11 | # A workflow run is made up of one or more jobs that can run sequentially or in parallel
12 | jobs:
13 | # This workflow contains a single job called "build"
14 | build:
15 | # The type of runner that the job will run on
16 | runs-on: ubuntu-22.04
17 | strategy:
18 | matrix:
19 | python-version: ["3.8", "3.9", "3.10"]
20 |
21 | # Steps represent a sequence of tasks that will be executed as part of the job
22 | steps:
23 | # Checks-out your repository under $GITHUB_WORKSPACE, so your job can access it
24 | - uses: actions/checkout@v2
25 |
26 | - name: Set up Python ${{ matrix.python-version }}
27 | uses: actions/setup-python@v1
28 | with:
29 | python-version: ${{ matrix.python-version }}
30 |
31 | - name: Install dependencies
32 | run: |
33 | python -m pip install --upgrade pip
34 | pip install -r requirements.txt
35 | pip install isort==4.3.21
36 | pip install flake8==3.8.3
37 | pip install "importlib-metadata<5.0"
38 | # Runs a set of commands using the runners shell
39 | - name: Format check
40 | run: ./.github/workflows/format_check.sh
41 |
--------------------------------------------------------------------------------
/yolox/utils/checkpoint.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii Inc. All rights reserved.
4 | import os
5 | import shutil
6 | from loguru import logger
7 |
8 | import torch
9 |
10 |
11 | def load_ckpt(model, ckpt):
12 | model_state_dict = model.state_dict()
13 | load_dict = {}
14 | for key_model, v in model_state_dict.items():
15 | if key_model not in ckpt:
16 | logger.warning(
17 | "{} is not in the ckpt. Please double check and see if this is desired.".format(
18 | key_model
19 | )
20 | )
21 | continue
22 | v_ckpt = ckpt[key_model]
23 | if v.shape != v_ckpt.shape:
24 | logger.warning(
25 | "Shape of {} in checkpoint is {}, while shape of {} in model is {}.".format(
26 | key_model, v_ckpt.shape, key_model, v.shape
27 | )
28 | )
29 | continue
30 | load_dict[key_model] = v_ckpt
31 |
32 | model.load_state_dict(load_dict, strict=False)
33 | return model
34 |
35 |
36 | def save_checkpoint(state, is_best, save_dir, model_name=""):
37 | if not os.path.exists(save_dir):
38 | os.makedirs(save_dir)
39 | filename = os.path.join(save_dir, model_name + "_ckpt.pth")
40 | torch.save(state, filename)
41 | if is_best:
42 | best_filename = os.path.join(save_dir, "best_ckpt.pth")
43 | shutil.copyfile(filename, best_filename)
44 |
--------------------------------------------------------------------------------
/demo/MegEngine/python/dump.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
5 | import argparse
6 |
7 | import megengine as mge
8 | import numpy as np
9 | from megengine import jit
10 |
11 | from build import build_and_load
12 |
13 |
14 | def make_parser():
15 | parser = argparse.ArgumentParser("YOLOX Demo Dump")
16 | parser.add_argument("-n", "--name", type=str, default="yolox-s", help="model name")
17 | parser.add_argument("-c", "--ckpt", default=None, type=str, help="ckpt for eval")
18 | parser.add_argument(
19 | "--dump_path", default="model.mge", help="path to save the dumped model"
20 | )
21 | return parser
22 |
23 |
24 | def dump_static_graph(model, graph_name="model.mge"):
25 | model.eval()
26 | model.head.decode_in_inference = False
27 |
28 | data = mge.Tensor(np.random.random((1, 3, 640, 640)))
29 |
30 | @jit.trace(capture_as_const=True)
31 | def pred_func(data):
32 | outputs = model(data)
33 | return outputs
34 |
35 | pred_func(data)
36 | pred_func.dump(
37 | graph_name,
38 | arg_names=["data"],
39 | optimize_for_inference=True,
40 | enable_fuse_conv_bias_nonlinearity=True,
41 | )
42 |
43 |
44 | def main(args):
45 | model = build_and_load(args.ckpt, name=args.name)
46 | dump_static_graph(model, args.dump_path)
47 |
48 |
49 | if __name__ == "__main__":
50 | args = make_parser().parse_args()
51 | main(args)
52 |
--------------------------------------------------------------------------------
/SECURITY.md:
--------------------------------------------------------------------------------
1 | # Security Policy
2 |
3 | ## Reporting a Vulnerability
4 |
5 | ### Types of Security Issues
6 | We actively monitor:
7 | - Code vulnerabilities (RCE, XSS, authentication bypass)
8 | - Dependency risks (critical vulnerabilities in project dependencies, such as requirements.txt, pyproject.toml, or equivalent files)
9 | - Configuration flaws (insecure defaults in deployment scripts)
10 |
11 | ### Disclosure Channels (Choose one):
12 |
13 | 1. **Encrypted Email**
14 | Contact: `wangfeng19950315@163.com`
15 | *Subject format: `[SECURITY] ModuleName - Brief Description`*
16 |
17 | 2. **GitHub Private Report**
18 | Use GitHub's ["Report a vulnerability"](https://github.com/Megvii-BaseDetection/YOLOX/security/advisories) feature
19 |
20 | 3. **Reporting Security Issues**
21 | Please report security issues using Create new issue: https://github.com/Megvii-BaseDetection/YOLOX/issues/new
22 |
23 |
24 | ## Response Process
25 | 1. **Acknowledgement**
26 | - Initial response within **48 business hours**
27 | 2. **Assessment**
28 | - Triage using CVSS v3.1 scoring
29 | 3. **Remediation**
30 | - Critical (CVSS ≥9.0): Patch within **7 days**
31 | - High (CVSS 7-8.9): Patch within **30 days**
32 | 4. **Public Disclosure**
33 | - Published via [GitHub Advisories](https://github.com/Megvii-BaseDetection/YOLOX/security/advisories)
34 | - CVE assignment coordinated with [MITRE](https://cveform.mitre.org)
35 |
36 | ## Secure Development Practices
37 | - Always verify hashes when downloading dependencies:
38 | ```bash
39 | sha256sum -c
40 | ```
41 |
--------------------------------------------------------------------------------
/docs/assignment_visualization.md:
--------------------------------------------------------------------------------
1 | # Visualize label assignment
2 |
3 | This tutorial explains how to visualize your label asssignment result when training with YOLOX.
4 |
5 | ## 1. Visualization command
6 |
7 | We provide a visualization tool to help you visualize your label assignment result. You can find it in [`tools/visualize_assignment.py`](../tools/visualize_assign.py).
8 |
9 | Here is an example of command to visualize your label assignment result:
10 |
11 | ```shell
12 | python3 tools/visualize_assign.py -f /path/to/your/exp.py yolox-s -d 1 -b 8 --max-batch 2
13 | ```
14 |
15 | `max-batch` here means the maximum number of batches to visualize. The default value is 1, which the tool means only visualize the first batch.
16 |
17 | By the way, the mosaic augmentation is used in default dataloader, so you can also see the mosaic result here.
18 |
19 | After running the command, the logger will show you where the visualization result is saved, let's open it and into the step 2.
20 |
21 | ## 2. Check the visualization result
22 |
23 | Here is an example of visualization result:
24 | 
25 |
26 | Those dots in one box is the matched anchor of gt box. **The color of dots is the same as the color of the box** to help you determine which object is assigned to the anchor. Note the box and dots are **instance level** visualization, which means the same class may have different colors.
27 | **If the gt box doesn't match any anchor, the box will be marked as red and the red text "unmatched" will be drawn over the box**.
28 |
29 | Please feel free to open an issue if you have any questions.
30 |
--------------------------------------------------------------------------------
/demo/nebullvm/nebullvm_optimization.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import time
3 | from nebullvm.api.functions import optimize_model # Install DL compilers
4 | from yolox.exp import get_exp
5 |
6 | # Get YOLO model
7 | exp = get_exp(None, 'yolox-s') # select model name
8 | model = exp.get_model()
9 | model.cuda()
10 | model.eval()
11 |
12 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
13 |
14 | # Create dummy data for the optimizer
15 | input_data = [((torch.randn(1, 3, 640, 640).to(device), ), 0) for i in range(100)]
16 |
17 | # ---------- Optimization ----------
18 | optimized_model = optimize_model(model, input_data=input_data, optimization_time="constrained") # Optimization without performance loss
19 |
20 |
21 | # ---------- Benchmarks ----------
22 | # Select image to test the latency of the optimized model
23 |
24 | # Create dummy image
25 | img = torch.randn(1, 3, 640, 640).to(device)
26 |
27 | # Check perfomance
28 | warmup_iters = 30
29 | num_iters = 100
30 |
31 | # Unptimized model perfomance
32 | with torch.no_grad():
33 | for i in range(warmup_iters):
34 | o = model(img)
35 |
36 | start = time.time()
37 | for i in range(num_iters):
38 | o = model(img)
39 | stop = time.time()
40 | print(f"Average inference time of unoptimized YOLOX: {(stop - start)/num_iters*1000} ms")
41 |
42 | # Optimized model perfomance
43 | with torch.no_grad():
44 | for i in range(warmup_iters):
45 | res = optimized_model(img)
46 |
47 | start = time.time()
48 | for i in range(num_iters):
49 | res = optimized_model(img)
50 | stop = time.time()
51 | print(f"Average inference time of YOLOX otpimized with nebullvm: {(stop - start)/num_iters*1000} ms")
52 |
--------------------------------------------------------------------------------
/yolox/data/datasets/coco_classes.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
5 | COCO_CLASSES = (
6 | "person",
7 | "bicycle",
8 | "car",
9 | "motorcycle",
10 | "airplane",
11 | "bus",
12 | "train",
13 | "truck",
14 | "boat",
15 | "traffic light",
16 | "fire hydrant",
17 | "stop sign",
18 | "parking meter",
19 | "bench",
20 | "bird",
21 | "cat",
22 | "dog",
23 | "horse",
24 | "sheep",
25 | "cow",
26 | "elephant",
27 | "bear",
28 | "zebra",
29 | "giraffe",
30 | "backpack",
31 | "umbrella",
32 | "handbag",
33 | "tie",
34 | "suitcase",
35 | "frisbee",
36 | "skis",
37 | "snowboard",
38 | "sports ball",
39 | "kite",
40 | "baseball bat",
41 | "baseball glove",
42 | "skateboard",
43 | "surfboard",
44 | "tennis racket",
45 | "bottle",
46 | "wine glass",
47 | "cup",
48 | "fork",
49 | "knife",
50 | "spoon",
51 | "bowl",
52 | "banana",
53 | "apple",
54 | "sandwich",
55 | "orange",
56 | "broccoli",
57 | "carrot",
58 | "hot dog",
59 | "pizza",
60 | "donut",
61 | "cake",
62 | "chair",
63 | "couch",
64 | "potted plant",
65 | "bed",
66 | "dining table",
67 | "toilet",
68 | "tv",
69 | "laptop",
70 | "mouse",
71 | "remote",
72 | "keyboard",
73 | "cell phone",
74 | "microwave",
75 | "oven",
76 | "toaster",
77 | "sink",
78 | "refrigerator",
79 | "book",
80 | "clock",
81 | "vase",
82 | "scissors",
83 | "teddy bear",
84 | "hair drier",
85 | "toothbrush",
86 | )
87 |
--------------------------------------------------------------------------------
/demo/TensorRT/python/README.md:
--------------------------------------------------------------------------------
1 | # YOLOX-TensorRT in Python
2 |
3 | This tutorial includes a Python demo for TensorRT.
4 |
5 | ## Install TensorRT Toolkit
6 |
7 | Please follow the [TensorRT Installation Guide](https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html) and [torch2trt gitrepo](https://github.com/NVIDIA-AI-IOT/torch2trt) to install TensorRT and torch2trt.
8 |
9 | ## Convert model
10 |
11 | YOLOX models can be easily conveted to TensorRT models using torch2trt
12 |
13 | If you want to convert our model, use the flag -n to specify a model name:
14 | ```shell
15 | python tools/trt.py -n -c
16 | ```
17 | For example:
18 | ```shell
19 | python tools/trt.py -n yolox-s -c your_ckpt.pth
20 | ```
21 | can be: yolox-nano, yolox-tiny. yolox-s, yolox-m, yolox-l, yolox-x.
22 |
23 | If you want to convert your customized model, use the flag -f to specify you exp file:
24 | ```shell
25 | python tools/trt.py -f -c
26 | ```
27 | For example:
28 | ```shell
29 | python tools/trt.py -f /path/to/your/yolox/exps/yolox_s.py -c your_ckpt.pth
30 | ```
31 | *yolox_s.py* can be any exp file modified by you.
32 |
33 | The converted model and the serialized engine file (for C++ demo) will be saved on your experiment output dir.
34 |
35 | ## Demo
36 |
37 | The TensorRT python demo is merged on our pytorch demo file, so you can run the pytorch demo command with ```--trt```.
38 |
39 | ```shell
40 | python tools/demo.py image -n yolox-s --trt --save_result
41 | ```
42 | or
43 | ```shell
44 | python tools/demo.py image -f exps/default/yolox_s.py --trt --save_result
45 | ```
46 |
47 |
--------------------------------------------------------------------------------
/exps/default/yolox_nano.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
5 | import os
6 |
7 | import torch.nn as nn
8 |
9 | from yolox.exp import Exp as MyExp
10 |
11 |
12 | class Exp(MyExp):
13 | def __init__(self):
14 | super(Exp, self).__init__()
15 | self.depth = 0.33
16 | self.width = 0.25
17 | self.input_size = (416, 416)
18 | self.random_size = (10, 20)
19 | self.mosaic_scale = (0.5, 1.5)
20 | self.test_size = (416, 416)
21 | self.mosaic_prob = 0.5
22 | self.enable_mixup = False
23 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
24 |
25 | def get_model(self, sublinear=False):
26 |
27 | def init_yolo(M):
28 | for m in M.modules():
29 | if isinstance(m, nn.BatchNorm2d):
30 | m.eps = 1e-3
31 | m.momentum = 0.03
32 | if "model" not in self.__dict__:
33 | from yolox.models import YOLOX, YOLOPAFPN, YOLOXHead
34 | in_channels = [256, 512, 1024]
35 | # NANO model use depthwise = True, which is main difference.
36 | backbone = YOLOPAFPN(
37 | self.depth, self.width, in_channels=in_channels,
38 | act=self.act, depthwise=True,
39 | )
40 | head = YOLOXHead(
41 | self.num_classes, self.width, in_channels=in_channels,
42 | act=self.act, depthwise=True
43 | )
44 | self.model = YOLOX(backbone, head)
45 |
46 | self.model.apply(init_yolo)
47 | self.model.head.initialize_biases(1e-2)
48 | return self.model
49 |
--------------------------------------------------------------------------------
/demo/MegEngine/python/build.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 |
4 | import megengine as mge
5 | import megengine.module as M
6 |
7 | from models.yolo_fpn import YOLOFPN
8 | from models.yolo_head import YOLOXHead
9 | from models.yolo_pafpn import YOLOPAFPN
10 | from models.yolox import YOLOX
11 |
12 |
13 | def build_yolox(name="yolox-s"):
14 | num_classes = 80
15 |
16 | # value meaning: depth, width
17 | param_dict = {
18 | "yolox-nano": (0.33, 0.25),
19 | "yolox-tiny": (0.33, 0.375),
20 | "yolox-s": (0.33, 0.50),
21 | "yolox-m": (0.67, 0.75),
22 | "yolox-l": (1.0, 1.0),
23 | "yolox-x": (1.33, 1.25),
24 | }
25 | if name == "yolov3":
26 | depth = 1.0
27 | width = 1.0
28 | backbone = YOLOFPN()
29 | head = YOLOXHead(num_classes, width, in_channels=[128, 256, 512], act="lrelu")
30 | model = YOLOX(backbone, head)
31 | else:
32 | assert name in param_dict
33 | kwargs = {}
34 | depth, width = param_dict[name]
35 | if name == "yolox-nano":
36 | kwargs["depthwise"] = True
37 | in_channels = [256, 512, 1024]
38 | backbone = YOLOPAFPN(depth, width, in_channels=in_channels, **kwargs)
39 | head = YOLOXHead(num_classes, width, in_channels=in_channels, **kwargs)
40 | model = YOLOX(backbone, head)
41 |
42 | for m in model.modules():
43 | if isinstance(m, M.BatchNorm2d):
44 | m.eps = 1e-3
45 |
46 | return model
47 |
48 |
49 | def build_and_load(weight_file, name="yolox-s"):
50 | model = build_yolox(name)
51 | model_weights = mge.load(weight_file)
52 | model.load_state_dict(model_weights, strict=False)
53 | return model
54 |
--------------------------------------------------------------------------------
/yolox/models/yolox.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python
2 | # -*- encoding: utf-8 -*-
3 | # Copyright (c) Megvii Inc. All rights reserved.
4 |
5 | import torch.nn as nn
6 |
7 | from .yolo_head import YOLOXHead
8 | from .yolo_pafpn import YOLOPAFPN
9 |
10 |
11 | class YOLOX(nn.Module):
12 | """
13 | YOLOX model module. The module list is defined by create_yolov3_modules function.
14 | The network returns loss values from three YOLO layers during training
15 | and detection results during test.
16 | """
17 |
18 | def __init__(self, backbone=None, head=None):
19 | super().__init__()
20 | if backbone is None:
21 | backbone = YOLOPAFPN()
22 | if head is None:
23 | head = YOLOXHead(80)
24 |
25 | self.backbone = backbone
26 | self.head = head
27 |
28 | def forward(self, x, targets=None):
29 | # fpn output content features of [dark3, dark4, dark5]
30 | fpn_outs = self.backbone(x)
31 |
32 | if self.training:
33 | assert targets is not None
34 | loss, iou_loss, conf_loss, cls_loss, l1_loss, num_fg = self.head(
35 | fpn_outs, targets, x
36 | )
37 | outputs = {
38 | "total_loss": loss,
39 | "iou_loss": iou_loss,
40 | "l1_loss": l1_loss,
41 | "conf_loss": conf_loss,
42 | "cls_loss": cls_loss,
43 | "num_fg": num_fg,
44 | }
45 | else:
46 | outputs = self.head(fpn_outs)
47 |
48 | return outputs
49 |
50 | def visualize(self, x, targets, save_prefix="assign_vis_"):
51 | fpn_outs = self.backbone(x)
52 | self.head.visualize_assign_result(fpn_outs, targets, x, save_prefix)
53 |
--------------------------------------------------------------------------------
/exps/example/custom/nano.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
5 | import os
6 |
7 | import torch.nn as nn
8 |
9 | from yolox.exp import Exp as MyExp
10 |
11 |
12 | class Exp(MyExp):
13 | def __init__(self):
14 | super(Exp, self).__init__()
15 | self.depth = 0.33
16 | self.width = 0.25
17 | self.input_size = (416, 416)
18 | self.mosaic_scale = (0.5, 1.5)
19 | self.random_size = (10, 20)
20 | self.test_size = (416, 416)
21 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
22 | self.enable_mixup = False
23 |
24 | # Define yourself dataset path
25 | self.data_dir = "datasets/coco128"
26 | self.train_ann = "instances_train2017.json"
27 | self.val_ann = "instances_val2017.json"
28 |
29 | self.num_classes = 71
30 |
31 | def get_model(self, sublinear=False):
32 |
33 | def init_yolo(M):
34 | for m in M.modules():
35 | if isinstance(m, nn.BatchNorm2d):
36 | m.eps = 1e-3
37 | m.momentum = 0.03
38 | if "model" not in self.__dict__:
39 | from yolox.models import YOLOX, YOLOPAFPN, YOLOXHead
40 | in_channels = [256, 512, 1024]
41 | # NANO model use depthwise = True, which is main difference.
42 | backbone = YOLOPAFPN(self.depth, self.width, in_channels=in_channels, depthwise=True)
43 | head = YOLOXHead(self.num_classes, self.width, in_channels=in_channels, depthwise=True)
44 | self.model = YOLOX(backbone, head)
45 |
46 | self.model.apply(init_yolo)
47 | self.model.head.initialize_biases(1e-2)
48 | return self.model
49 |
--------------------------------------------------------------------------------
/demo/TensorRT/cpp/README.md:
--------------------------------------------------------------------------------
1 | # YOLOX-TensorRT in C++
2 |
3 | As YOLOX models are easy to convert to tensorrt using [torch2trt gitrepo](https://github.com/NVIDIA-AI-IOT/torch2trt),
4 | our C++ demo does not include the model converting or constructing like other tenorrt demos.
5 |
6 |
7 | ## Step 1: Prepare serialized engine file
8 |
9 | Follow the trt [python demo README](https://github.com/Megvii-BaseDetection/YOLOX/blob/main/demo/TensorRT/python/README.md) to convert and save the serialized engine file.
10 |
11 | Check the 'model_trt.engine' file generated from Step 1, which will be automatically saved at the current demo dir.
12 |
13 |
14 | ## Step 2: build the demo
15 |
16 | Please follow the [TensorRT Installation Guide](https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html) to install TensorRT.
17 |
18 | And you should set the TensorRT path and CUDA path in CMakeLists.txt.
19 |
20 | If you train your custom dataset, you may need to modify the value of `num_class`.
21 |
22 | ```c++
23 | const int num_class = 80;
24 | ```
25 |
26 | Install opencv with ```sudo apt-get install libopencv-dev``` (we don't need a higher version of opencv like v3.3+).
27 |
28 | build the demo:
29 |
30 | ```shell
31 | mkdir build
32 | cd build
33 | cmake ..
34 | make
35 | ```
36 |
37 | Then run the demo:
38 |
39 | ```shell
40 | ./yolox ../model_trt.engine -i ../../../../assets/dog.jpg
41 | ```
42 |
43 | or
44 |
45 | ```shell
46 | ./yolox -i
47 | ```
48 |
49 | NOTE: for `trtexec` users, modify `INPUT_BLOB_NAME` and `OUTPUT_BLOB_NAME` as the following code.
50 | ```
51 | const char* INPUT_BLOB_NAME = "images";
52 | const char* OUTPUT_BLOB_NAME = "output";
53 | ```
54 |
55 | Here is the command to convert the small onnx model to tensorrt engine file:
56 | ```
57 | trtexec --onnx=yolox_s.onnx --saveEngine=yolox_s.trt
58 | ```
59 |
--------------------------------------------------------------------------------
/yolox/data/data_prefetcher.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
5 | import torch
6 |
7 |
8 | class DataPrefetcher:
9 | """
10 | DataPrefetcher is inspired by code of following file:
11 | https://github.com/NVIDIA/apex/blob/master/examples/imagenet/main_amp.py
12 | It could speedup your pytorch dataloader. For more information, please check
13 | https://github.com/NVIDIA/apex/issues/304#issuecomment-493562789.
14 | """
15 |
16 | def __init__(self, loader):
17 | self.loader = iter(loader)
18 | self.stream = torch.cuda.Stream()
19 | self.input_cuda = self._input_cuda_for_image
20 | self.record_stream = DataPrefetcher._record_stream_for_image
21 | self.preload()
22 |
23 | def preload(self):
24 | try:
25 | self.next_input, self.next_target, _, _ = next(self.loader)
26 | except StopIteration:
27 | self.next_input = None
28 | self.next_target = None
29 | return
30 |
31 | with torch.cuda.stream(self.stream):
32 | self.input_cuda()
33 | self.next_target = self.next_target.cuda(non_blocking=True)
34 |
35 | def next(self):
36 | torch.cuda.current_stream().wait_stream(self.stream)
37 | input = self.next_input
38 | target = self.next_target
39 | if input is not None:
40 | self.record_stream(input)
41 | if target is not None:
42 | target.record_stream(torch.cuda.current_stream())
43 | self.preload()
44 | return input, target
45 |
46 | def _input_cuda_for_image(self):
47 | self.next_input = self.next_input.cuda(non_blocking=True)
48 |
49 | @staticmethod
50 | def _record_stream_for_image(input):
51 | input.record_stream(torch.cuda.current_stream())
52 |
--------------------------------------------------------------------------------
/yolox/models/losses.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python
2 | # -*- encoding: utf-8 -*-
3 | # Copyright (c) Megvii Inc. All rights reserved.
4 |
5 | import torch
6 | import torch.nn as nn
7 |
8 |
9 | class IOUloss(nn.Module):
10 | def __init__(self, reduction="none", loss_type="iou"):
11 | super(IOUloss, self).__init__()
12 | self.reduction = reduction
13 | self.loss_type = loss_type
14 |
15 | def forward(self, pred, target):
16 | assert pred.shape[0] == target.shape[0]
17 |
18 | pred = pred.view(-1, 4)
19 | target = target.view(-1, 4)
20 | tl = torch.max(
21 | (pred[:, :2] - pred[:, 2:] / 2), (target[:, :2] - target[:, 2:] / 2)
22 | )
23 | br = torch.min(
24 | (pred[:, :2] + pred[:, 2:] / 2), (target[:, :2] + target[:, 2:] / 2)
25 | )
26 |
27 | area_p = torch.prod(pred[:, 2:], 1)
28 | area_g = torch.prod(target[:, 2:], 1)
29 |
30 | en = (tl < br).type(tl.type()).prod(dim=1)
31 | area_i = torch.prod(br - tl, 1) * en
32 | area_u = area_p + area_g - area_i
33 | iou = (area_i) / (area_u + 1e-16)
34 |
35 | if self.loss_type == "iou":
36 | loss = 1 - iou ** 2
37 | elif self.loss_type == "giou":
38 | c_tl = torch.min(
39 | (pred[:, :2] - pred[:, 2:] / 2), (target[:, :2] - target[:, 2:] / 2)
40 | )
41 | c_br = torch.max(
42 | (pred[:, :2] + pred[:, 2:] / 2), (target[:, :2] + target[:, 2:] / 2)
43 | )
44 | area_c = torch.prod(c_br - c_tl, 1)
45 | giou = iou - (area_c - area_u) / area_c.clamp(1e-16)
46 | loss = 1 - giou.clamp(min=-1.0, max=1.0)
47 |
48 | if self.reduction == "mean":
49 | loss = loss.mean()
50 | elif self.reduction == "sum":
51 | loss = loss.sum()
52 |
53 | return loss
54 |
--------------------------------------------------------------------------------
/demo/MegEngine/python/convert_weights.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | import argparse
4 | from collections import OrderedDict
5 |
6 | import megengine as mge
7 | import torch
8 |
9 |
10 | def make_parser():
11 | parser = argparse.ArgumentParser()
12 | parser.add_argument("-w", "--weights", type=str, help="path of weight file")
13 | parser.add_argument(
14 | "-o",
15 | "--output",
16 | default="weight_mge.pkl",
17 | type=str,
18 | help="path of weight file",
19 | )
20 | return parser
21 |
22 |
23 | def numpy_weights(weight_file):
24 | torch_weights = torch.load(weight_file, map_location="cpu")
25 | if "model" in torch_weights:
26 | torch_weights = torch_weights["model"]
27 | new_dict = OrderedDict()
28 | for k, v in torch_weights.items():
29 | new_dict[k] = v.cpu().numpy()
30 | return new_dict
31 |
32 |
33 | def map_weights(weight_file, output_file):
34 | torch_weights = numpy_weights(weight_file)
35 |
36 | new_dict = OrderedDict()
37 | for k, v in torch_weights.items():
38 | if "num_batches_tracked" in k:
39 | print("drop: {}".format(k))
40 | continue
41 | if k.endswith("bias"):
42 | print("bias key: {}".format(k))
43 | v = v.reshape(1, -1, 1, 1)
44 | new_dict[k] = v
45 | elif "dconv" in k and "conv.weight" in k:
46 | print("depthwise conv key: {}".format(k))
47 | cout, cin, k1, k2 = v.shape
48 | v = v.reshape(cout, 1, cin, k1, k2)
49 | new_dict[k] = v
50 | else:
51 | new_dict[k] = v
52 |
53 | mge.save(new_dict, output_file)
54 | print("save weights to {}".format(output_file))
55 |
56 |
57 | def main():
58 | parser = make_parser()
59 | args = parser.parse_args()
60 | map_weights(args.weights, args.output)
61 |
62 |
63 | if __name__ == "__main__":
64 | main()
65 |
--------------------------------------------------------------------------------
/docs/updates_note.md:
--------------------------------------------------------------------------------
1 |
2 | # Updates notes
3 |
4 | ## 【2021/08/19】
5 |
6 | * Support image caching for faster training, which requires large system RAM.
7 | * Remove the dependence of apex and support torch amp training.
8 | * Optimize the preprocessing for faster training
9 | * Replace the older distort augmentation with new HSV aug for faster training and better performance.
10 |
11 | ### 2X Faster training
12 |
13 | We optimize the data preprocess and support image caching with `--cache` flag:
14 |
15 | ```shell
16 | python tools/train.py -n yolox-s -d 8 -b 64 --fp16 -o [--cache]
17 | yolox-m
18 | yolox-l
19 | yolox-x
20 | ```
21 | * -d: number of gpu devices
22 | * -b: total batch size, the recommended number for -b is num-gpu * 8
23 | * --fp16: mixed precision training
24 | * --cache: caching imgs into RAM to accelarate training, which need large system RAM.
25 |
26 | ### Higher performance
27 |
28 | New models achieve **~1%** higher performance! See [Model_Zoo](model_zoo.md) for more details.
29 |
30 | ### Support torch amp
31 |
32 | We now support torch.cuda.amp training and Apex is not used anymore.
33 |
34 | ### Breaking changes
35 |
36 | We remove the normalization operation like -mean/std. This will make the old weights **incompatible**.
37 |
38 | If you still want to use old weights, you can add `--legacy' in demo and eval:
39 |
40 | ```shell
41 | python tools/demo.py image -n yolox-s -c /path/to/your/yolox_s.pth --path assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result --device [cpu/gpu] [--legacy]
42 | ```
43 |
44 | and
45 |
46 | ```shell
47 | python tools/eval.py -n yolox-s -c yolox_s.pth -b 64 -d 8 --conf 0.001 [--fp16] [--fuse] [--legacy]
48 | yolox-m
49 | yolox-l
50 | yolox-x
51 | ```
52 |
53 | But for deployment demo, we don't support the old weights anymore. Users could checkout to YOLOX version 0.1.0 to use legacy weights for deployment
54 |
55 |
56 |
--------------------------------------------------------------------------------
/docs/manipulate_training_image_size.md:
--------------------------------------------------------------------------------
1 | # Manipulating Your Training Image Size
2 |
3 | This tutorial explains how to control your image size when training on your own data.
4 |
5 | ## 1. Introduction
6 |
7 | There are 3 hyperparamters control the training size:
8 |
9 | - self.input_size = (640, 640) #(height, width)
10 | - self.multiscale_range = 5
11 | - self.random_size = (14, 26)
12 |
13 | There is 1 hyperparameter constrols the testing size:
14 |
15 | - self.test_size = (640, 640)
16 |
17 | The self.input_size is suggested to set to the same value as self.test_size. By default, it is set to (640, 640) for most models and (416, 416) for yolox-tiny and yolox-nano.
18 |
19 | ## 2. Multi Scale Training
20 |
21 | When training on your custom dataset, you can use multiscale training in 2 ways:
22 |
23 | 1. **【Default】Only specifying the self.input_size and leaving others unchanged.**
24 |
25 | If so, the actual multiscale sizes range from:
26 |
27 | [self.input_size[0] - self.multiscale_range\*32, self.input_size[0] + self.multiscale_range\*32]
28 |
29 | For example, if you only set:
30 |
31 | ```python
32 | self.input_size = (640, 640)
33 | ```
34 |
35 | the actual multiscale range is [640 - 5*32, 640 + 5\*32], i.e., [480, 800].
36 |
37 | You can modify self.multiscale_range to change the multiscale range.
38 |
39 | 2. **Simultaneously specifying the self.input_size and self.random_size**
40 |
41 | ```python
42 | self.input_size = (416, 416)
43 | self.random_size = (10, 20)
44 | ```
45 |
46 | In this case, the actual multiscale range is [self.random_size[0]\*32, self.random_size[1]\*32], i.e., [320, 640]
47 |
48 | **Note: You must specify the self.input_size because it is used for initializing resize aug in dataset.**
49 |
50 | ## 3. Single Scale Training
51 |
52 | If you want to train in a single scale. You need to specify the self.input_size and self.multiscale_range=0:
53 |
54 | ```python
55 | self.input_size = (416, 416)
56 | self.multiscale_range = 0
57 | ```
58 |
59 | **DO NOT** set the self.random_size.
60 |
--------------------------------------------------------------------------------
/exps/example/yolox_voc/yolox_voc_s.py:
--------------------------------------------------------------------------------
1 | # encoding: utf-8
2 | import os
3 |
4 | from yolox.data import get_yolox_datadir
5 | from yolox.exp import Exp as MyExp
6 |
7 |
8 | class Exp(MyExp):
9 | def __init__(self):
10 | super(Exp, self).__init__()
11 | self.num_classes = 20
12 | self.depth = 0.33
13 | self.width = 0.50
14 | self.warmup_epochs = 1
15 |
16 | # ---------- transform config ------------ #
17 | self.mosaic_prob = 1.0
18 | self.mixup_prob = 1.0
19 | self.hsv_prob = 1.0
20 | self.flip_prob = 0.5
21 |
22 | self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
23 |
24 | def get_dataset(self, cache: bool, cache_type: str = "ram"):
25 | from yolox.data import VOCDetection, TrainTransform
26 |
27 | return VOCDetection(
28 | data_dir=os.path.join(get_yolox_datadir(), "VOCdevkit"),
29 | image_sets=[('2007', 'trainval'), ('2012', 'trainval')],
30 | img_size=self.input_size,
31 | preproc=TrainTransform(
32 | max_labels=50,
33 | flip_prob=self.flip_prob,
34 | hsv_prob=self.hsv_prob),
35 | cache=cache,
36 | cache_type=cache_type,
37 | )
38 |
39 | def get_eval_dataset(self, **kwargs):
40 | from yolox.data import VOCDetection, ValTransform
41 | legacy = kwargs.get("legacy", False)
42 |
43 | return VOCDetection(
44 | data_dir=os.path.join(get_yolox_datadir(), "VOCdevkit"),
45 | image_sets=[('2007', 'test')],
46 | img_size=self.test_size,
47 | preproc=ValTransform(legacy=legacy),
48 | )
49 |
50 | def get_evaluator(self, batch_size, is_distributed, testdev=False, legacy=False):
51 | from yolox.evaluators import VOCEvaluator
52 |
53 | return VOCEvaluator(
54 | dataloader=self.get_eval_loader(batch_size, is_distributed,
55 | testdev=testdev, legacy=legacy),
56 | img_size=self.test_size,
57 | confthre=self.test_conf,
58 | nmsthre=self.nmsthre,
59 | num_classes=self.num_classes,
60 | )
61 |
--------------------------------------------------------------------------------
/yolox/utils/ema.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii Inc. All rights reserved.
4 | import math
5 | from copy import deepcopy
6 |
7 | import torch
8 | import torch.nn as nn
9 |
10 | __all__ = ["ModelEMA", "is_parallel"]
11 |
12 |
13 | def is_parallel(model):
14 | """check if model is in parallel mode."""
15 | parallel_type = (
16 | nn.parallel.DataParallel,
17 | nn.parallel.DistributedDataParallel,
18 | )
19 | return isinstance(model, parallel_type)
20 |
21 |
22 | class ModelEMA:
23 | """
24 | Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models
25 | Keep a moving average of everything in the model state_dict (parameters and buffers).
26 | This is intended to allow functionality like
27 | https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage
28 | A smoothed version of the weights is necessary for some training schemes to perform well.
29 | This class is sensitive where it is initialized in the sequence of model init,
30 | GPU assignment and distributed training wrappers.
31 | """
32 |
33 | def __init__(self, model, decay=0.9999, updates=0):
34 | """
35 | Args:
36 | model (nn.Module): model to apply EMA.
37 | decay (float): ema decay reate.
38 | updates (int): counter of EMA updates.
39 | """
40 | # Create EMA(FP32)
41 | self.ema = deepcopy(model.module if is_parallel(model) else model).eval()
42 | self.updates = updates
43 | # decay exponential ramp (to help early epochs)
44 | self.decay = lambda x: decay * (1 - math.exp(-x / 2000))
45 | for p in self.ema.parameters():
46 | p.requires_grad_(False)
47 |
48 | def update(self, model):
49 | # Update EMA parameters
50 | with torch.no_grad():
51 | self.updates += 1
52 | d = self.decay(self.updates)
53 |
54 | msd = (
55 | model.module.state_dict() if is_parallel(model) else model.state_dict()
56 | ) # model state_dict
57 | for k, v in self.ema.state_dict().items():
58 | if v.dtype.is_floating_point:
59 | v *= d
60 | v += (1.0 - d) * msd[k].detach()
61 |
--------------------------------------------------------------------------------
/demo/ncnn/android/gradlew.bat:
--------------------------------------------------------------------------------
1 | @if "%DEBUG%" == "" @echo off
2 | @rem ##########################################################################
3 | @rem
4 | @rem Gradle startup script for Windows
5 | @rem
6 | @rem ##########################################################################
7 |
8 | @rem Set local scope for the variables with windows NT shell
9 | if "%OS%"=="Windows_NT" setlocal
10 |
11 | set DIRNAME=%~dp0
12 | if "%DIRNAME%" == "" set DIRNAME=.
13 | set APP_BASE_NAME=%~n0
14 | set APP_HOME=%DIRNAME%
15 |
16 | @rem Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script.
17 | set DEFAULT_JVM_OPTS=
18 |
19 | @rem Find java.exe
20 | if defined JAVA_HOME goto findJavaFromJavaHome
21 |
22 | set JAVA_EXE=java.exe
23 | %JAVA_EXE% -version >NUL 2>&1
24 | if "%ERRORLEVEL%" == "0" goto init
25 |
26 | echo.
27 | echo ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH.
28 | echo.
29 | echo Please set the JAVA_HOME variable in your environment to match the
30 | echo location of your Java installation.
31 |
32 | goto fail
33 |
34 | :findJavaFromJavaHome
35 | set JAVA_HOME=%JAVA_HOME:"=%
36 | set JAVA_EXE=%JAVA_HOME%/bin/java.exe
37 |
38 | if exist "%JAVA_EXE%" goto init
39 |
40 | echo.
41 | echo ERROR: JAVA_HOME is set to an invalid directory: %JAVA_HOME%
42 | echo.
43 | echo Please set the JAVA_HOME variable in your environment to match the
44 | echo location of your Java installation.
45 |
46 | goto fail
47 |
48 | :init
49 | @rem Get command-line arguments, handling Windows variants
50 |
51 | if not "%OS%" == "Windows_NT" goto win9xME_args
52 |
53 | :win9xME_args
54 | @rem Slurp the command line arguments.
55 | set CMD_LINE_ARGS=
56 | set _SKIP=2
57 |
58 | :win9xME_args_slurp
59 | if "x%~1" == "x" goto execute
60 |
61 | set CMD_LINE_ARGS=%*
62 |
63 | :execute
64 | @rem Setup the command line
65 |
66 | set CLASSPATH=%APP_HOME%\gradle\wrapper\gradle-wrapper.jar
67 |
68 | @rem Execute Gradle
69 | "%JAVA_EXE%" %DEFAULT_JVM_OPTS% %JAVA_OPTS% %GRADLE_OPTS% "-Dorg.gradle.appname=%APP_BASE_NAME%" -classpath "%CLASSPATH%" org.gradle.wrapper.GradleWrapperMain %CMD_LINE_ARGS%
70 |
71 | :end
72 | @rem End local scope for the variables with windows NT shell
73 | if "%ERRORLEVEL%"=="0" goto mainEnd
74 |
75 | :fail
76 | rem Set variable GRADLE_EXIT_CONSOLE if you need the _script_ return code instead of
77 | rem the _cmd.exe /c_ return code!
78 | if not "" == "%GRADLE_EXIT_CONSOLE%" exit 1
79 | exit /b 1
80 |
81 | :mainEnd
82 | if "%OS%"=="Windows_NT" endlocal
83 |
84 | :omega
85 |
--------------------------------------------------------------------------------
/demo/MegEngine/python/models/yolo_fpn.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- encoding: utf-8 -*-
3 | # Copyright (c) Megvii Inc. All rights reserved.
4 |
5 | import megengine.functional as F
6 | import megengine.module as M
7 |
8 | from .darknet import Darknet
9 | from .network_blocks import BaseConv, UpSample
10 |
11 |
12 | class YOLOFPN(M.Module):
13 | """
14 | YOLOFPN module. Darknet 53 is the default backbone of this model.
15 | """
16 |
17 | def __init__(
18 | self, depth=53, in_features=["dark3", "dark4", "dark5"],
19 | ):
20 | super().__init__()
21 |
22 | self.backbone = Darknet(depth)
23 | self.in_features = in_features
24 |
25 | # out 1
26 | self.out1_cbl = self._make_cbl(512, 256, 1)
27 | self.out1 = self._make_embedding([256, 512], 512 + 256)
28 |
29 | # out 2
30 | self.out2_cbl = self._make_cbl(256, 128, 1)
31 | self.out2 = self._make_embedding([128, 256], 256 + 128)
32 |
33 | # upsample
34 | self.upsample = UpSample(scale_factor=2, mode="bilinear")
35 |
36 | def _make_cbl(self, _in, _out, ks):
37 | return BaseConv(_in, _out, ks, stride=1, act="lrelu")
38 |
39 | def _make_embedding(self, filters_list, in_filters):
40 | m = M.Sequential(
41 | *[
42 | self._make_cbl(in_filters, filters_list[0], 1),
43 | self._make_cbl(filters_list[0], filters_list[1], 3),
44 |
45 | self._make_cbl(filters_list[1], filters_list[0], 1),
46 |
47 | self._make_cbl(filters_list[0], filters_list[1], 3),
48 | self._make_cbl(filters_list[1], filters_list[0], 1),
49 | ]
50 | )
51 | return m
52 |
53 | def forward(self, inputs):
54 | """
55 | Args:
56 | inputs (Tensor): input image.
57 |
58 | Returns:
59 | Tuple[Tensor]: FPN output features..
60 | """
61 | # backbone
62 | out_features = self.backbone(inputs)
63 | x2, x1, x0 = [out_features[f] for f in self.in_features]
64 |
65 | # yolo branch 1
66 | x1_in = self.out1_cbl(x0)
67 | x1_in = self.upsample(x1_in)
68 | x1_in = F.concat([x1_in, x1], 1)
69 | out_dark4 = self.out1(x1_in)
70 |
71 | # yolo branch 2
72 | x2_in = self.out2_cbl(out_dark4)
73 | x2_in = self.upsample(x2_in)
74 | x2_in = F.concat([x2_in, x2], 1)
75 | out_dark3 = self.out2(x2_in)
76 |
77 | outputs = (out_dark3, out_dark4, x0)
78 | return outputs
79 |
--------------------------------------------------------------------------------
/tools/export_torchscript.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
5 | import argparse
6 | import os
7 | from loguru import logger
8 |
9 | import torch
10 |
11 | from yolox.exp import get_exp
12 |
13 |
14 | def make_parser():
15 | parser = argparse.ArgumentParser("YOLOX torchscript deploy")
16 | parser.add_argument(
17 | "--output-name", type=str, default="yolox.torchscript.pt", help="output name of models"
18 | )
19 | parser.add_argument("--batch-size", type=int, default=1, help="batch size")
20 | parser.add_argument(
21 | "-f",
22 | "--exp_file",
23 | default=None,
24 | type=str,
25 | help="experiment description file",
26 | )
27 | parser.add_argument("-expn", "--experiment-name", type=str, default=None)
28 | parser.add_argument("-n", "--name", type=str, default=None, help="model name")
29 | parser.add_argument("-c", "--ckpt", default=None, type=str, help="ckpt path")
30 | parser.add_argument(
31 | "--decode_in_inference",
32 | action="store_true",
33 | help="decode in inference or not"
34 | )
35 | parser.add_argument(
36 | "opts",
37 | help="Modify config options using the command-line",
38 | default=None,
39 | nargs=argparse.REMAINDER,
40 | )
41 |
42 | return parser
43 |
44 |
45 | @logger.catch
46 | def main():
47 | args = make_parser().parse_args()
48 | logger.info("args value: {}".format(args))
49 | exp = get_exp(args.exp_file, args.name)
50 | exp.merge(args.opts)
51 |
52 | if not args.experiment_name:
53 | args.experiment_name = exp.exp_name
54 |
55 | model = exp.get_model()
56 | if args.ckpt is None:
57 | file_name = os.path.join(exp.output_dir, args.experiment_name)
58 | ckpt_file = os.path.join(file_name, "best_ckpt.pth")
59 | else:
60 | ckpt_file = args.ckpt
61 |
62 | # load the model state dict
63 | ckpt = torch.load(ckpt_file, map_location="cpu")
64 |
65 | model.eval()
66 | if "model" in ckpt:
67 | ckpt = ckpt["model"]
68 | model.load_state_dict(ckpt)
69 | model.head.decode_in_inference = args.decode_in_inference
70 |
71 | logger.info("loading checkpoint done.")
72 | dummy_input = torch.randn(args.batch_size, 3, exp.test_size[0], exp.test_size[1])
73 |
74 | mod = torch.jit.trace(model, dummy_input)
75 | mod.save(args.output_name)
76 | logger.info("generated torchscript model named {}".format(args.output_name))
77 |
78 |
79 | if __name__ == "__main__":
80 | main()
81 |
--------------------------------------------------------------------------------
/demo/MegEngine/cpp/build.sh:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env bash
2 | set -e
3 |
4 | if [ -z $CXX ];then
5 | echo "please export you c++ toolchain to CXX"
6 | echo "for example:"
7 | echo "build for host: export CXX=g++"
8 | echo "cross build for aarch64-android(always locate in NDK): export CXX=aarch64-linux-android21-clang++"
9 | echo "cross build for aarch64-linux: export CXX=aarch64-linux-gnu-g++"
10 | exit -1
11 | fi
12 |
13 | if [ -z $MGE_INSTALL_PATH ];then
14 | echo "please refsi ./README.md to init MGE_INSTALL_PATH env"
15 | exit -1
16 | fi
17 |
18 | if [ -z $OPENCV_INSTALL_INCLUDE_PATH ];then
19 | echo "please refs ./README.md to init OPENCV_INSTALL_INCLUDE_PATH env"
20 | exit -1
21 | fi
22 |
23 | if [ -z $OPENCV_INSTALL_LIB_PATH ];then
24 | echo "please refs ./README.md to init OPENCV_INSTALL_LIB_PATH env"
25 | exit -1
26 | fi
27 |
28 | INCLUDE_FLAG="-I$MGE_INSTALL_PATH/include -I$OPENCV_INSTALL_INCLUDE_PATH"
29 | LINK_FLAG="-L$MGE_INSTALL_PATH/lib/ -lmegengine -L$OPENCV_INSTALL_LIB_PATH -lopencv_core -lopencv_highgui -lopencv_imgproc -lopencv_imgcodecs"
30 | BUILD_FLAG="-static-libstdc++ -O3 -pie -fPIE -g"
31 |
32 | if [[ $CXX =~ "android" ]]; then
33 | LINK_FLAG="${LINK_FLAG} -llog -lz"
34 | fi
35 |
36 | echo "CXX: $CXX"
37 | echo "MGE_INSTALL_PATH: $MGE_INSTALL_PATH"
38 | echo "INCLUDE_FLAG: $INCLUDE_FLAG"
39 | echo "LINK_FLAG: $LINK_FLAG"
40 | echo "BUILD_FLAG: $BUILD_FLAG"
41 |
42 | echo "[" > compile_commands.json
43 | echo "{" >> compile_commands.json
44 | echo "\"directory\": \"$PWD\"," >> compile_commands.json
45 | echo "\"command\": \"$CXX yolox.cpp -o yolox ${INCLUDE_FLAG} ${LINK_FLAG}\"," >> compile_commands.json
46 | echo "\"file\": \"$PWD/yolox.cpp\"," >> compile_commands.json
47 | echo "}," >> compile_commands.json
48 | echo "]" >> compile_commands.json
49 | $CXX yolox.cpp -o yolox ${INCLUDE_FLAG} ${LINK_FLAG} ${BUILD_FLAG}
50 |
51 | echo "build success, output file: yolox"
52 | if [[ $CXX =~ "android" ]]; then
53 | echo "try command to run:"
54 | echo "adb push/scp $MGE_INSTALL_PATH/lib/libmegengine.so android_phone"
55 | echo "adb push/scp $OPENCV_INSTALL_LIB_PATH/*.so android_phone"
56 | echo "adb push/scp ./yolox yolox_s.mge android_phone"
57 | echo "adb push/scp ../../../assets/dog.jpg android_phone"
58 | echo "adb/ssh to android_phone, then run: LD_LIBRARY_PATH=. ./yolox yolox_s.mge dog.jpg cpu/multithread "
59 | else
60 | echo "try command to run: LD_LIBRARY_PATH=$MGE_INSTALL_PATH/lib/:$OPENCV_INSTALL_LIB_PATH ./yolox yolox_s.mge ../../../assets/dog.jpg cuda/cpu/multithread "
61 | fi
62 |
--------------------------------------------------------------------------------
/yolox/models/yolo_fpn.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python
2 | # -*- encoding: utf-8 -*-
3 | # Copyright (c) Megvii Inc. All rights reserved.
4 |
5 | import torch
6 | import torch.nn as nn
7 |
8 | from .darknet import Darknet
9 | from .network_blocks import BaseConv
10 |
11 |
12 | class YOLOFPN(nn.Module):
13 | """
14 | YOLOFPN module. Darknet 53 is the default backbone of this model.
15 | """
16 |
17 | def __init__(
18 | self,
19 | depth=53,
20 | in_features=["dark3", "dark4", "dark5"],
21 | ):
22 | super().__init__()
23 |
24 | self.backbone = Darknet(depth)
25 | self.in_features = in_features
26 |
27 | # out 1
28 | self.out1_cbl = self._make_cbl(512, 256, 1)
29 | self.out1 = self._make_embedding([256, 512], 512 + 256)
30 |
31 | # out 2
32 | self.out2_cbl = self._make_cbl(256, 128, 1)
33 | self.out2 = self._make_embedding([128, 256], 256 + 128)
34 |
35 | # upsample
36 | self.upsample = nn.Upsample(scale_factor=2, mode="nearest")
37 |
38 | def _make_cbl(self, _in, _out, ks):
39 | return BaseConv(_in, _out, ks, stride=1, act="lrelu")
40 |
41 | def _make_embedding(self, filters_list, in_filters):
42 | m = nn.Sequential(
43 | *[
44 | self._make_cbl(in_filters, filters_list[0], 1),
45 | self._make_cbl(filters_list[0], filters_list[1], 3),
46 | self._make_cbl(filters_list[1], filters_list[0], 1),
47 | self._make_cbl(filters_list[0], filters_list[1], 3),
48 | self._make_cbl(filters_list[1], filters_list[0], 1),
49 | ]
50 | )
51 | return m
52 |
53 | def load_pretrained_model(self, filename="./weights/darknet53.mix.pth"):
54 | with open(filename, "rb") as f:
55 | state_dict = torch.load(f, map_location="cpu")
56 | print("loading pretrained weights...")
57 | self.backbone.load_state_dict(state_dict)
58 |
59 | def forward(self, inputs):
60 | """
61 | Args:
62 | inputs (Tensor): input image.
63 |
64 | Returns:
65 | Tuple[Tensor]: FPN output features..
66 | """
67 | # backbone
68 | out_features = self.backbone(inputs)
69 | x2, x1, x0 = [out_features[f] for f in self.in_features]
70 |
71 | # yolo branch 1
72 | x1_in = self.out1_cbl(x0)
73 | x1_in = self.upsample(x1_in)
74 | x1_in = torch.cat([x1_in, x1], 1)
75 | out_dark4 = self.out1(x1_in)
76 |
77 | # yolo branch 2
78 | x2_in = self.out2_cbl(out_dark4)
79 | x2_in = self.upsample(x2_in)
80 | x2_in = torch.cat([x2_in, x2], 1)
81 | out_dark3 = self.out2(x2_in)
82 |
83 | outputs = (out_dark3, out_dark4, x0)
84 | return outputs
85 |
--------------------------------------------------------------------------------
/tools/trt.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
5 | import argparse
6 | import os
7 | import shutil
8 | from loguru import logger
9 |
10 | import tensorrt as trt
11 | import torch
12 | from torch2trt import torch2trt
13 |
14 | from yolox.exp import get_exp
15 |
16 |
17 | def make_parser():
18 | parser = argparse.ArgumentParser("YOLOX ncnn deploy")
19 | parser.add_argument("-expn", "--experiment-name", type=str, default=None)
20 | parser.add_argument("-n", "--name", type=str, default=None, help="model name")
21 |
22 | parser.add_argument(
23 | "-f",
24 | "--exp_file",
25 | default=None,
26 | type=str,
27 | help="please input your experiment description file",
28 | )
29 | parser.add_argument("-c", "--ckpt", default=None, type=str, help="ckpt path")
30 | parser.add_argument(
31 | "-w", '--workspace', type=int, default=32, help='max workspace size in detect'
32 | )
33 | parser.add_argument("-b", '--batch', type=int, default=1, help='max batch size in detect')
34 | return parser
35 |
36 |
37 | @logger.catch
38 | @torch.no_grad()
39 | def main():
40 | args = make_parser().parse_args()
41 | exp = get_exp(args.exp_file, args.name)
42 | if not args.experiment_name:
43 | args.experiment_name = exp.exp_name
44 |
45 | model = exp.get_model()
46 | file_name = os.path.join(exp.output_dir, args.experiment_name)
47 | os.makedirs(file_name, exist_ok=True)
48 | if args.ckpt is None:
49 | ckpt_file = os.path.join(file_name, "best_ckpt.pth")
50 | else:
51 | ckpt_file = args.ckpt
52 |
53 | ckpt = torch.load(ckpt_file, map_location="cpu")
54 | # load the model state dict
55 |
56 | model.load_state_dict(ckpt["model"])
57 | logger.info("loaded checkpoint done.")
58 | model.eval()
59 | model.cuda()
60 | model.head.decode_in_inference = False
61 | x = torch.ones(1, 3, exp.test_size[0], exp.test_size[1]).cuda()
62 | model_trt = torch2trt(
63 | model,
64 | [x],
65 | fp16_mode=True,
66 | log_level=trt.Logger.INFO,
67 | max_workspace_size=(1 << args.workspace),
68 | max_batch_size=args.batch,
69 | )
70 | torch.save(model_trt.state_dict(), os.path.join(file_name, "model_trt.pth"))
71 | logger.info("Converted TensorRT model done.")
72 | engine_file = os.path.join(file_name, "model_trt.engine")
73 | engine_file_demo = os.path.join("demo", "TensorRT", "cpp", "model_trt.engine")
74 | with open(engine_file, "wb") as f:
75 | f.write(model_trt.engine.serialize())
76 |
77 | shutil.copyfile(engine_file, engine_file_demo)
78 |
79 | logger.info("Converted TensorRT model engine file is saved for C++ inference.")
80 |
81 |
82 | if __name__ == "__main__":
83 | main()
84 |
--------------------------------------------------------------------------------
/demo/ONNXRuntime/onnx_inference.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # Copyright (c) Megvii, Inc. and its affiliates.
3 |
4 | import argparse
5 | import os
6 |
7 | import cv2
8 | import numpy as np
9 |
10 | import onnxruntime
11 |
12 | from yolox.data.data_augment import preproc as preprocess
13 | from yolox.data.datasets import COCO_CLASSES
14 | from yolox.utils import mkdir, multiclass_nms, demo_postprocess, vis
15 |
16 |
17 | def make_parser():
18 | parser = argparse.ArgumentParser("onnxruntime inference sample")
19 | parser.add_argument(
20 | "-m",
21 | "--model",
22 | type=str,
23 | default="yolox.onnx",
24 | help="Input your onnx model.",
25 | )
26 | parser.add_argument(
27 | "-i",
28 | "--image_path",
29 | type=str,
30 | default='test_image.png',
31 | help="Path to your input image.",
32 | )
33 | parser.add_argument(
34 | "-o",
35 | "--output_dir",
36 | type=str,
37 | default='demo_output',
38 | help="Path to your output directory.",
39 | )
40 | parser.add_argument(
41 | "-s",
42 | "--score_thr",
43 | type=float,
44 | default=0.3,
45 | help="Score threshould to filter the result.",
46 | )
47 | parser.add_argument(
48 | "--input_shape",
49 | type=str,
50 | default="640,640",
51 | help="Specify an input shape for inference.",
52 | )
53 | return parser
54 |
55 |
56 | if __name__ == '__main__':
57 | args = make_parser().parse_args()
58 |
59 | input_shape = tuple(map(int, args.input_shape.split(',')))
60 | origin_img = cv2.imread(args.image_path)
61 | img, ratio = preprocess(origin_img, input_shape)
62 |
63 | session = onnxruntime.InferenceSession(args.model)
64 |
65 | ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}
66 | output = session.run(None, ort_inputs)
67 | predictions = demo_postprocess(output[0], input_shape)[0]
68 |
69 | boxes = predictions[:, :4]
70 | scores = predictions[:, 4:5] * predictions[:, 5:]
71 |
72 | boxes_xyxy = np.ones_like(boxes)
73 | boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2.
74 | boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2.
75 | boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2.
76 | boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2.
77 | boxes_xyxy /= ratio
78 | dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1)
79 | if dets is not None:
80 | final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
81 | origin_img = vis(origin_img, final_boxes, final_scores, final_cls_inds,
82 | conf=args.score_thr, class_names=COCO_CLASSES)
83 |
84 | mkdir(args.output_dir)
85 | output_path = os.path.join(args.output_dir, os.path.basename(args.image_path))
86 | cv2.imwrite(output_path, origin_img)
87 |
--------------------------------------------------------------------------------
/yolox/exp/base_exp.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # Copyright (c) Megvii Inc. All rights reserved.
3 |
4 | import ast
5 | import pprint
6 | from abc import ABCMeta, abstractmethod
7 | from typing import Dict, List, Tuple
8 | from tabulate import tabulate
9 |
10 | import torch
11 | from torch.nn import Module
12 |
13 | from yolox.utils import LRScheduler
14 |
15 |
16 | class BaseExp(metaclass=ABCMeta):
17 | """Basic class for any experiment."""
18 |
19 | def __init__(self):
20 | self.seed = None
21 | self.output_dir = "./YOLOX_outputs"
22 | self.print_interval = 100
23 | self.eval_interval = 10
24 | self.dataset = None
25 |
26 | @abstractmethod
27 | def get_model(self) -> Module:
28 | pass
29 |
30 | @abstractmethod
31 | def get_dataset(self, cache: bool = False, cache_type: str = "ram"):
32 | pass
33 |
34 | @abstractmethod
35 | def get_data_loader(
36 | self, batch_size: int, is_distributed: bool
37 | ) -> Dict[str, torch.utils.data.DataLoader]:
38 | pass
39 |
40 | @abstractmethod
41 | def get_optimizer(self, batch_size: int) -> torch.optim.Optimizer:
42 | pass
43 |
44 | @abstractmethod
45 | def get_lr_scheduler(
46 | self, lr: float, iters_per_epoch: int, **kwargs
47 | ) -> LRScheduler:
48 | pass
49 |
50 | @abstractmethod
51 | def get_evaluator(self):
52 | pass
53 |
54 | @abstractmethod
55 | def eval(self, model, evaluator, weights):
56 | pass
57 |
58 | def __repr__(self):
59 | table_header = ["keys", "values"]
60 | exp_table = [
61 | (str(k), pprint.pformat(v))
62 | for k, v in vars(self).items()
63 | if not k.startswith("_")
64 | ]
65 | return tabulate(exp_table, headers=table_header, tablefmt="fancy_grid")
66 |
67 | def merge(self, cfg_list):
68 | assert len(cfg_list) % 2 == 0, f"length must be even, check value here: {cfg_list}"
69 | for k, v in zip(cfg_list[0::2], cfg_list[1::2]):
70 | # only update value with same key
71 | if hasattr(self, k):
72 | src_value = getattr(self, k)
73 | src_type = type(src_value)
74 |
75 | # pre-process input if source type is list or tuple
76 | if isinstance(src_value, (List, Tuple)):
77 | v = v.strip("[]()")
78 | v = [t.strip() for t in v.split(",")]
79 |
80 | # find type of tuple
81 | if len(src_value) > 0:
82 | src_item_type = type(src_value[0])
83 | v = [src_item_type(t) for t in v]
84 |
85 | if src_value is not None and src_type != type(v):
86 | try:
87 | v = src_type(v)
88 | except Exception:
89 | v = ast.literal_eval(v)
90 | setattr(self, k, v)
91 |
--------------------------------------------------------------------------------
/yolox/utils/setup_env.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii Inc. All rights reserved.
4 |
5 | import os
6 | import subprocess
7 | from loguru import logger
8 |
9 | import cv2
10 |
11 | from .dist import get_world_size, is_main_process
12 |
13 | __all__ = ["configure_nccl", "configure_module", "configure_omp"]
14 |
15 |
16 | def configure_nccl():
17 | """Configure multi-machine environment variables of NCCL."""
18 | os.environ["NCCL_LAUNCH_MODE"] = "PARALLEL"
19 | os.environ["NCCL_IB_HCA"] = subprocess.getoutput(
20 | "pushd /sys/class/infiniband/ > /dev/null; for i in mlx5_*; "
21 | "do cat $i/ports/1/gid_attrs/types/* 2>/dev/null "
22 | "| grep v >/dev/null && echo $i ; done; popd > /dev/null"
23 | )
24 | os.environ["NCCL_IB_GID_INDEX"] = "3"
25 | os.environ["NCCL_IB_TC"] = "106"
26 |
27 |
28 | def configure_omp(num_threads=1):
29 | """
30 | If OMP_NUM_THREADS is not configured and world_size is greater than 1,
31 | Configure OMP_NUM_THREADS environment variables of NCCL to `num_thread`.
32 |
33 | Args:
34 | num_threads (int): value of `OMP_NUM_THREADS` to set.
35 | """
36 | # We set OMP_NUM_THREADS=1 by default, which achieves the best speed on our machines
37 | # feel free to change it for better performance.
38 | if "OMP_NUM_THREADS" not in os.environ and get_world_size() > 1:
39 | os.environ["OMP_NUM_THREADS"] = str(num_threads)
40 | if is_main_process():
41 | logger.info(
42 | "\n***************************************************************\n"
43 | "We set `OMP_NUM_THREADS` for each process to {} to speed up.\n"
44 | "please further tune the variable for optimal performance.\n"
45 | "***************************************************************".format(
46 | os.environ["OMP_NUM_THREADS"]
47 | )
48 | )
49 |
50 |
51 | def configure_module(ulimit_value=8192):
52 | """
53 | Configure pytorch module environment. setting of ulimit and cv2 will be set.
54 |
55 | Args:
56 | ulimit_value(int): default open file number on linux. Default value: 8192.
57 | """
58 | # system setting
59 | try:
60 | import resource
61 |
62 | rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
63 | resource.setrlimit(resource.RLIMIT_NOFILE, (ulimit_value, rlimit[1]))
64 | except Exception:
65 | # Exception might be raised in Windows OS or rlimit reaches max limit number.
66 | # However, set rlimit value might not be necessary.
67 | pass
68 |
69 | # cv2
70 | # multiprocess might be harmful on performance of torch dataloader
71 | os.environ["OPENCV_OPENCL_RUNTIME"] = "disabled"
72 | try:
73 | cv2.setNumThreads(0)
74 | cv2.ocl.setUseOpenCL(False)
75 | except Exception:
76 | # cv2 version mismatch might rasie exceptions.
77 | pass
78 |
--------------------------------------------------------------------------------
/setup.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python
2 | # Copyright (c) Megvii, Inc. and its affiliates. All Rights Reserved
3 |
4 | import re
5 | import setuptools
6 | import sys
7 |
8 | TORCH_AVAILABLE = True
9 | try:
10 | import torch
11 | from torch.utils import cpp_extension
12 | except ImportError:
13 | TORCH_AVAILABLE = False
14 | print("[WARNING] Unable to import torch, pre-compiling ops will be disabled.")
15 |
16 |
17 | def get_package_dir():
18 | pkg_dir = {
19 | "yolox.tools": "tools",
20 | "yolox.exp.default": "exps/default",
21 | }
22 | return pkg_dir
23 |
24 |
25 | def get_install_requirements():
26 | with open("requirements.txt", "r", encoding="utf-8") as f:
27 | reqs = [x.strip() for x in f.read().splitlines()]
28 | reqs = [x for x in reqs if not x.startswith("#")]
29 | return reqs
30 |
31 |
32 | def get_yolox_version():
33 | with open("yolox/__init__.py", "r") as f:
34 | version = re.search(
35 | r'^__version__\s*=\s*[\'"]([^\'"]*)[\'"]',
36 | f.read(), re.MULTILINE
37 | ).group(1)
38 | return version
39 |
40 |
41 | def get_long_description():
42 | with open("README.md", "r", encoding="utf-8") as f:
43 | long_description = f.read()
44 | return long_description
45 |
46 |
47 | def get_ext_modules():
48 | ext_module = []
49 | if sys.platform != "win32": # pre-compile ops on linux
50 | assert TORCH_AVAILABLE, "torch is required for pre-compiling ops, please install it first."
51 | # if any other op is added, please also add it here
52 | from yolox.layers import FastCOCOEvalOp
53 | ext_module.append(FastCOCOEvalOp().build_op())
54 | return ext_module
55 |
56 |
57 | def get_cmd_class():
58 | cmdclass = {}
59 | if TORCH_AVAILABLE:
60 | cmdclass["build_ext"] = cpp_extension.BuildExtension
61 | return cmdclass
62 |
63 |
64 | setuptools.setup(
65 | name="yolox",
66 | version=get_yolox_version(),
67 | author="megvii basedet team",
68 | url="https://github.com/Megvii-BaseDetection/YOLOX",
69 | package_dir=get_package_dir(),
70 | packages=setuptools.find_packages(exclude=("tests", "tools")) + list(get_package_dir().keys()),
71 | python_requires=">=3.6",
72 | install_requires=get_install_requirements(),
73 | setup_requires=["wheel"], # avoid building error when pip is not updated
74 | long_description=get_long_description(),
75 | long_description_content_type="text/markdown",
76 | include_package_data=True, # include files in MANIFEST.in
77 | ext_modules=get_ext_modules(),
78 | cmdclass=get_cmd_class(),
79 | classifiers=[
80 | "Programming Language :: Python :: 3", "Operating System :: OS Independent",
81 | "License :: OSI Approved :: Apache Software License",
82 | ],
83 | project_urls={
84 | "Documentation": "https://yolox.readthedocs.io",
85 | "Source": "https://github.com/Megvii-BaseDetection/YOLOX",
86 | "Tracker": "https://github.com/Megvii-BaseDetection/YOLOX/issues",
87 | },
88 | )
89 |
--------------------------------------------------------------------------------
/tools/visualize_assign.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # Copyright (c) Megvii, Inc. and its affiliates.
3 |
4 | import os
5 | import sys
6 | import random
7 | import time
8 | import warnings
9 | from loguru import logger
10 |
11 | import torch
12 | import torch.backends.cudnn as cudnn
13 |
14 | from yolox.exp import Exp, get_exp
15 | from yolox.core import Trainer
16 | from yolox.utils import configure_module, configure_omp
17 | from yolox.tools.train import make_parser
18 |
19 |
20 | class AssignVisualizer(Trainer):
21 |
22 | def __init__(self, exp: Exp, args):
23 | super().__init__(exp, args)
24 | self.batch_cnt = 0
25 | self.vis_dir = os.path.join(self.file_name, "vis")
26 | os.makedirs(self.vis_dir, exist_ok=True)
27 |
28 | def train_one_iter(self):
29 | iter_start_time = time.time()
30 |
31 | inps, targets = self.prefetcher.next()
32 | inps = inps.to(self.data_type)
33 | targets = targets.to(self.data_type)
34 | targets.requires_grad = False
35 | inps, targets = self.exp.preprocess(inps, targets, self.input_size)
36 | data_end_time = time.time()
37 |
38 | with torch.cuda.amp.autocast(enabled=self.amp_training):
39 | path_prefix = os.path.join(self.vis_dir, f"assign_vis_{self.batch_cnt}_")
40 | self.model.visualize(inps, targets, path_prefix)
41 |
42 | if self.use_model_ema:
43 | self.ema_model.update(self.model)
44 |
45 | iter_end_time = time.time()
46 | self.meter.update(
47 | iter_time=iter_end_time - iter_start_time,
48 | data_time=data_end_time - iter_start_time,
49 | )
50 | self.batch_cnt += 1
51 | if self.batch_cnt >= self.args.max_batch:
52 | sys.exit(0)
53 |
54 | def after_train(self):
55 | logger.info("Finish visualize assignment, exit...")
56 |
57 |
58 | def assign_vis_parser():
59 | parser = make_parser()
60 | parser.add_argument("--max-batch", type=int, default=1, help="max batch of images to visualize")
61 | return parser
62 |
63 |
64 | @logger.catch
65 | def main(exp: Exp, args):
66 | if exp.seed is not None:
67 | random.seed(exp.seed)
68 | torch.manual_seed(exp.seed)
69 | cudnn.deterministic = True
70 | warnings.warn(
71 | "You have chosen to seed training. This will turn on the CUDNN deterministic setting, "
72 | "which can slow down your training considerably! You may see unexpected behavior "
73 | "when restarting from checkpoints."
74 | )
75 |
76 | # set environment variables for distributed training
77 | configure_omp()
78 | cudnn.benchmark = True
79 |
80 | visualizer = AssignVisualizer(exp, args)
81 | visualizer.train()
82 |
83 |
84 | if __name__ == "__main__":
85 | configure_module()
86 | args = assign_vis_parser().parse_args()
87 | exp = get_exp(args.exp_file, args.name)
88 | exp.merge(args.opts)
89 |
90 | if not args.experiment_name:
91 | args.experiment_name = exp.exp_name
92 |
93 | main(exp, args)
94 |
--------------------------------------------------------------------------------
/yolox/data/samplers.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
5 | import itertools
6 | from typing import Optional
7 |
8 | import torch
9 | import torch.distributed as dist
10 | from torch.utils.data.sampler import BatchSampler as torchBatchSampler
11 | from torch.utils.data.sampler import Sampler
12 |
13 |
14 | class YoloBatchSampler(torchBatchSampler):
15 | """
16 | This batch sampler will generate mini-batches of (mosaic, index) tuples from another sampler.
17 | It works just like the :class:`torch.utils.data.sampler.BatchSampler`,
18 | but it will turn on/off the mosaic aug.
19 | """
20 |
21 | def __init__(self, *args, mosaic=True, **kwargs):
22 | super().__init__(*args, **kwargs)
23 | self.mosaic = mosaic
24 |
25 | def __iter__(self):
26 | for batch in super().__iter__():
27 | yield [(self.mosaic, idx) for idx in batch]
28 |
29 |
30 | class InfiniteSampler(Sampler):
31 | """
32 | In training, we only care about the "infinite stream" of training data.
33 | So this sampler produces an infinite stream of indices and
34 | all workers cooperate to correctly shuffle the indices and sample different indices.
35 | The samplers in each worker effectively produces `indices[worker_id::num_workers]`
36 | where `indices` is an infinite stream of indices consisting of
37 | `shuffle(range(size)) + shuffle(range(size)) + ...` (if shuffle is True)
38 | or `range(size) + range(size) + ...` (if shuffle is False)
39 | """
40 |
41 | def __init__(
42 | self,
43 | size: int,
44 | shuffle: bool = True,
45 | seed: Optional[int] = 0,
46 | rank=0,
47 | world_size=1,
48 | ):
49 | """
50 | Args:
51 | size (int): the total number of data of the underlying dataset to sample from
52 | shuffle (bool): whether to shuffle the indices or not
53 | seed (int): the initial seed of the shuffle. Must be the same
54 | across all workers. If None, will use a random seed shared
55 | among workers (require synchronization among all workers).
56 | """
57 | self._size = size
58 | assert size > 0
59 | self._shuffle = shuffle
60 | self._seed = int(seed)
61 |
62 | if dist.is_available() and dist.is_initialized():
63 | self._rank = dist.get_rank()
64 | self._world_size = dist.get_world_size()
65 | else:
66 | self._rank = rank
67 | self._world_size = world_size
68 |
69 | def __iter__(self):
70 | start = self._rank
71 | yield from itertools.islice(
72 | self._infinite_indices(), start, None, self._world_size
73 | )
74 |
75 | def _infinite_indices(self):
76 | g = torch.Generator()
77 | g.manual_seed(self._seed)
78 | while True:
79 | if self._shuffle:
80 | yield from torch.randperm(self._size, generator=g)
81 | else:
82 | yield from torch.arange(self._size)
83 |
84 | def __len__(self):
85 | return self._size // self._world_size
86 |
--------------------------------------------------------------------------------
/yolox/utils/allreduce_norm.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii Inc. All rights reserved.
4 |
5 | import pickle
6 | from collections import OrderedDict
7 |
8 | import torch
9 | from torch import distributed as dist
10 | from torch import nn
11 |
12 | from .dist import _get_global_gloo_group, get_world_size
13 |
14 | ASYNC_NORM = (
15 | nn.BatchNorm1d,
16 | nn.BatchNorm2d,
17 | nn.BatchNorm3d,
18 | nn.InstanceNorm1d,
19 | nn.InstanceNorm2d,
20 | nn.InstanceNorm3d,
21 | )
22 |
23 | __all__ = [
24 | "get_async_norm_states",
25 | "pyobj2tensor",
26 | "tensor2pyobj",
27 | "all_reduce",
28 | "all_reduce_norm",
29 | ]
30 |
31 |
32 | def get_async_norm_states(module):
33 | async_norm_states = OrderedDict()
34 | for name, child in module.named_modules():
35 | if isinstance(child, ASYNC_NORM):
36 | for k, v in child.state_dict().items():
37 | async_norm_states[".".join([name, k])] = v
38 | return async_norm_states
39 |
40 |
41 | def pyobj2tensor(pyobj, device="cuda"):
42 | """serialize picklable python object to tensor"""
43 | storage = torch.ByteStorage.from_buffer(pickle.dumps(pyobj))
44 | return torch.ByteTensor(storage).to(device=device)
45 |
46 |
47 | def tensor2pyobj(tensor):
48 | """deserialize tensor to picklable python object"""
49 | return pickle.loads(tensor.cpu().numpy().tobytes())
50 |
51 |
52 | def _get_reduce_op(op_name):
53 | return {
54 | "sum": dist.ReduceOp.SUM,
55 | "mean": dist.ReduceOp.SUM,
56 | }[op_name.lower()]
57 |
58 |
59 | def all_reduce(py_dict, op="sum", group=None):
60 | """
61 | Apply all reduce function for python dict object.
62 | NOTE: make sure that every py_dict has the same keys and values are in the same shape.
63 |
64 | Args:
65 | py_dict (dict): dict to apply all reduce op.
66 | op (str): operator, could be "sum" or "mean".
67 | """
68 | world_size = get_world_size()
69 | if world_size == 1:
70 | return py_dict
71 | if group is None:
72 | group = _get_global_gloo_group()
73 | if dist.get_world_size(group) == 1:
74 | return py_dict
75 |
76 | # all reduce logic across different devices.
77 | py_key = list(py_dict.keys())
78 | py_key_tensor = pyobj2tensor(py_key)
79 | dist.broadcast(py_key_tensor, src=0)
80 | py_key = tensor2pyobj(py_key_tensor)
81 |
82 | tensor_shapes = [py_dict[k].shape for k in py_key]
83 | tensor_numels = [py_dict[k].numel() for k in py_key]
84 |
85 | flatten_tensor = torch.cat([py_dict[k].flatten() for k in py_key])
86 | dist.all_reduce(flatten_tensor, op=_get_reduce_op(op))
87 | if op == "mean":
88 | flatten_tensor /= world_size
89 |
90 | split_tensors = [
91 | x.reshape(shape)
92 | for x, shape in zip(torch.split(flatten_tensor, tensor_numels), tensor_shapes)
93 | ]
94 | return OrderedDict({k: v for k, v in zip(py_key, split_tensors)})
95 |
96 |
97 | def all_reduce_norm(module):
98 | """
99 | All reduce norm statistics in different devices.
100 | """
101 | states = get_async_norm_states(module)
102 | states = all_reduce(states, op="mean")
103 | module.load_state_dict(states, strict=False)
104 |
--------------------------------------------------------------------------------
/docs/mlflow_integration.md:
--------------------------------------------------------------------------------
1 | ## MLFlow Integration
2 | YOLOX now supports MLFlow integration. MLFlow is an open-source platform for managing the end-to-end machine learning lifecycle. It is designed to work with any ML library, algorithm, deployment tool, or language. MLFlow can be used to track experiments, metrics, and parameters, and to log and visualize model artifacts. \
3 | For more information, please refer to: [MLFlow Documentation](https://www.mlflow.org/docs/latest/index.html)
4 |
5 | ## Follow these steps to start logging your experiments to MLFlow:
6 | ### Step-1: Install MLFlow via pip
7 | ```bash
8 | pip install mlflow python-dotenv
9 | ```
10 |
11 | ### Step-2: Set up MLFlow Tracking Server
12 | Start or connect to a MLFlow tracking server like databricks. You can start a local tracking server by running the following command:
13 | ```bash
14 | mlflow server --host 127.0.0.1 --port 8080
15 | ```
16 | Read more about setting up MLFlow tracking server [here](https://mlflow.org/docs/latest/tracking/server.html#mlflow-tracking-server)
17 |
18 | ### Step-3: Set up MLFlow Environment Variables
19 | Set the following environment variables in your `.env` file:
20 | ```bash
21 | MLFLOW_TRACKING_URI="127.0.0.1:5000" # set to your mlflow server URI
22 | MLFLOW_EXPERIMENT_NAME="/path/to/experiment" # set to your experiment name
23 | MLFLOW_TAGS={"release.candidate": "DEV1", "release.version": "0.0.0"}
24 | # config related to logging model to mlflow as pyfunc
25 | YOLOX_MLFLOW_LOG_MODEL_ARTIFACTS="True" # whether to log model (best or historical) or not
26 | YOLOX_MLFLOW_LOG_MODEL_PER_n_EPOCHS=30 # try logging model only after every n epochs
27 | YOLOX_MLFLOW_LOG_Nth_EPOCH_MODELS="False" # whether to log step model along with best_model or not
28 | YOLOX_MLFLOW_RUN_NAME="" # give a custom name to your run, otherwise a random name is assign by mlflow
29 | YOLOX_MLFLOW_FLATTEN_PARAMS="True" # flatten any sub sub params of dict to be logged as simple key value pair
30 |
31 |
32 | MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING=True # log system gpu usage and other metrices
33 | MLFLOW_NESTED_RUN="False" #whether to run as a nested run of given run_id
34 | MLFLOW_RUN_ID="" # continue training from a given run_id
35 | ```
36 | ### Step-5: Provide --logger "mlflow" to the training script
37 | ```bash
38 | python tools/train.py -l mlflow -f exps/path/to/exp.py -d 1 -b 8 --fp16 -o -c
39 | pre_trained_model/.pth
40 | # note the -l mlflow flag
41 | # one working example is this
42 | python tools/train.py -l mlflow -f exps/example/custom/yolox_s.py -d 1 -b 8 --fp16 -o -c pre_trained_model/yolox_s.pth
43 | ```
44 | ### Step-4: optional; start the mlflow ui and track your experiments
45 | If you log runs to a local mlruns directory, run the following command in the directory above it, then access http://127.0.0.1:5000 in your browser.
46 |
47 | ```bash
48 | mlflow ui --port 5000
49 | ```
50 |
51 | ## Optional Databricks Integration
52 |
53 | ### Step-1: Install Databricks sdk
54 | ```bash
55 | pip install databricks-sdk
56 | ```
57 |
58 | ### Step-2: Set up Databricks Environment Variables
59 | Set the following environment variables in your `.env` file:
60 | ```bash
61 | MLFLOW_TRACKING_URI="databricks" # set to databricks
62 | MLFLOW_EXPERIMENT_NAME="/Users///"
63 | DATABRICKS_HOST = "https://dbc-1234567890123456.cloud.databricks.com" # set to your server URI
64 | DATABRICKS_TOKEN = "dapixxxxxxxxxxxxx"
65 | ```
--------------------------------------------------------------------------------
/demo/ONNXRuntime/README.md:
--------------------------------------------------------------------------------
1 | ## YOLOX-ONNXRuntime in Python
2 |
3 | This doc introduces how to convert your pytorch model into onnx, and how to run an onnxruntime demo to verify your convertion.
4 |
5 | ### Step1: Install onnxruntime
6 |
7 | run the following command to install onnxruntime:
8 | ```shell
9 | pip install onnxruntime
10 | ```
11 |
12 | ### Step2: Get ONNX models
13 |
14 | Users might download our pre-generated ONNX models or convert their own models to ONNX.
15 |
16 | #### Download ONNX models.
17 |
18 | | Model | Parameters | GFLOPs | Test Size | mAP | Weights |
19 | |:------| :----: | :----: | :---: | :---: | :---: |
20 | | YOLOX-Nano | 0.91M | 1.08 | 416x416 | 25.8 |[github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_nano.onnx) |
21 | | YOLOX-Tiny | 5.06M | 6.45 | 416x416 |32.8 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_tiny.onnx) |
22 | | YOLOX-S | 9.0M | 26.8 | 640x640 |40.5 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_s.onnx) |
23 | | YOLOX-M | 25.3M | 73.8 | 640x640 |47.2 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_m.onnx) |
24 | | YOLOX-L | 54.2M | 155.6 | 640x640 |50.1 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_l.onnx) |
25 | | YOLOX-Darknet53| 63.72M | 185.3 | 640x640 |48.0 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_darknet.onnx) |
26 | | YOLOX-X | 99.1M | 281.9 | 640x640 |51.5 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_x.onnx) |
27 |
28 | #### Convert Your Model to ONNX
29 |
30 | First, you should move to by:
31 | ```shell
32 | cd
33 | ```
34 | Then, you can:
35 |
36 | 1. Convert a standard YOLOX model by -n:
37 | ```shell
38 | python3 tools/export_onnx.py --output-name yolox_s.onnx -n yolox-s -c yolox_s.pth
39 | ```
40 | Notes:
41 | * -n: specify a model name. The model name must be one of the [yolox-s,m,l,x and yolox-nano, yolox-tiny, yolov3]
42 | * -c: the model you have trained
43 | * -o: opset version, default 11. **However, if you will further convert your onnx model to [OpenVINO](https://github.com/Megvii-BaseDetection/YOLOX/demo/OpenVINO/), please specify the opset version to 10.**
44 | * --no-onnxsim: disable onnxsim
45 | * To customize an input shape for onnx model, modify the following code in tools/export.py:
46 |
47 | ```python
48 | dummy_input = torch.randn(1, 3, exp.test_size[0], exp.test_size[1])
49 | ```
50 |
51 | 1. Convert a standard YOLOX model by -f. When using -f, the above command is equivalent to:
52 |
53 | ```shell
54 | python3 tools/export_onnx.py --output-name yolox_s.onnx -f exps/default/yolox_s.py -c yolox_s.pth
55 | ```
56 |
57 | 3. To convert your customized model, please use -f:
58 |
59 | ```shell
60 | python3 tools/export_onnx.py --output-name your_yolox.onnx -f exps/your_dir/your_yolox.py -c your_yolox.pth
61 | ```
62 |
63 | ### Step3: ONNXRuntime Demo
64 |
65 | Step1.
66 | ```shell
67 | cd /demo/ONNXRuntime
68 | ```
69 |
70 | Step2.
71 | ```shell
72 | python3 onnx_inference.py -m -i -o -s 0.3 --input_shape 640,640
73 | ```
74 | Notes:
75 | * -m: your converted onnx model
76 | * -i: input_image
77 | * -s: score threshold for visualization.
78 | * --input_shape: should be consistent with the shape you used for onnx convertion.
79 |
--------------------------------------------------------------------------------
/demo/OpenVINO/python/README.md:
--------------------------------------------------------------------------------
1 | # YOLOX-OpenVINO in Python
2 |
3 | This tutorial includes a Python demo for OpenVINO, as well as some converted models.
4 |
5 | ### Download OpenVINO models.
6 |
7 | | Model | Parameters | GFLOPs | Test Size | mAP | Weights |
8 | |:------| :----: | :----: | :---: | :---: | :---: |
9 | | [YOLOX-Nano](../../../exps/default/nano.py) | 0.91M | 1.08 | 416x416 | 25.8 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_nano_openvino.tar.gz) |
10 | | [YOLOX-Tiny](../../../exps/default/yolox_tiny.py) | 5.06M | 6.45 | 416x416 |32.8 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_tiny_openvino.tar.gz) |
11 | | [YOLOX-S](../../../exps/default/yolox_s.py) | 9.0M | 26.8 | 640x640 |40.5 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_s_openvino.tar.gz) |
12 | | [YOLOX-M](../../../exps/default/yolox_m.py) | 25.3M | 73.8 | 640x640 |47.2 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_m_openvino.tar.gz) |
13 | | [YOLOX-L](../../../exps/default/yolox_l.py) | 54.2M | 155.6 | 640x640 |50.1 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_l_openvino.tar.gz) |
14 | | [YOLOX-Darknet53](../../../exps/default/yolov3.py) | 63.72M | 185.3 | 640x640 |48.0 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_dark_openvino.tar.gz) |
15 | | [YOLOX-X](../../../exps/default/yolox_x.py) | 99.1M | 281.9 | 640x640 |51.5 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_x_openvino.tar.gz) |
16 |
17 | ## Install OpenVINO Toolkit
18 |
19 | Please visit [Openvino Homepage](https://docs.openvinotoolkit.org/latest/get_started_guides.html) for more details.
20 |
21 | ## Set up the Environment
22 |
23 | ### For Linux
24 |
25 | **Option1. Set up the environment tempororally. You need to run this command everytime you start a new shell window.**
26 |
27 | ```shell
28 | source /opt/intel/openvino_2021/bin/setupvars.sh
29 | ```
30 |
31 | **Option2. Set up the environment permenantly.**
32 |
33 | *Step1.* For Linux:
34 | ```shell
35 | vim ~/.bashrc
36 | ```
37 |
38 | *Step2.* Add the following line into your file:
39 |
40 | ```shell
41 | source /opt/intel/openvino_2021/bin/setupvars.sh
42 | ```
43 |
44 | *Step3.* Save and exit the file, then run:
45 |
46 | ```shell
47 | source ~/.bashrc
48 | ```
49 |
50 |
51 | ## Convert model
52 |
53 | 1. Export ONNX model
54 |
55 | Please refer to the [ONNX tutorial](https://github.com/Megvii-BaseDetection/YOLOX/demo/ONNXRuntime). **Note that you should set --opset to 10, otherwise your next step will fail.**
56 |
57 | 2. Convert ONNX to OpenVINO
58 |
59 | ``` shell
60 | cd /openvino_2021/deployment_tools/model_optimizer
61 | ```
62 |
63 | Install requirements for convert tool
64 |
65 | ```shell
66 | sudo ./install_prerequisites/install_prerequisites_onnx.sh
67 | ```
68 |
69 | Then convert model.
70 | ```shell
71 | python3 mo.py --input_model --input_shape [--data_type FP16]
72 | ```
73 | For example:
74 | ```shell
75 | python3 mo.py --input_model yolox.onnx --input_shape [1,3,640,640] --data_type FP16 --output_dir converted_output
76 | ```
77 |
78 | ## Demo
79 |
80 | ### python
81 |
82 | ```shell
83 | python openvino_inference.py -m -i
84 | ```
85 | or
86 | ```shell
87 | python openvino_inference.py -m -i -o -s -d
88 | ```
89 |
90 |
--------------------------------------------------------------------------------
/demo/OpenVINO/cpp/README.md:
--------------------------------------------------------------------------------
1 | # YOLOX-OpenVINO in C++
2 |
3 | This tutorial includes a C++ demo for OpenVINO, as well as some converted models.
4 |
5 | ### Download OpenVINO models.
6 |
7 | | Model | Parameters | GFLOPs | Test Size | mAP | Weights |
8 | |:------| :----: | :----: | :---: | :---: | :---: |
9 | | [YOLOX-Nano](../../../exps/default/nano.py) | 0.91M | 1.08 | 416x416 | 25.8 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_nano_openvino.tar.gz) |
10 | | [YOLOX-Tiny](../../../exps/default/yolox_tiny.py) | 5.06M | 6.45 | 416x416 |32.8 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_tiny_openvino.tar.gz) |
11 | | [YOLOX-S](../../../exps/default/yolox_s.py) | 9.0M | 26.8 | 640x640 |40.5 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_s_openvino.tar.gz) |
12 | | [YOLOX-M](../../../exps/default/yolox_m.py) | 25.3M | 73.8 | 640x640 |47.2 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_m_openvino.tar.gz) |
13 | | [YOLOX-L](../../../exps/default/yolox_l.py) | 54.2M | 155.6 | 640x640 |50.1 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_l_openvino.tar.gz) |
14 | | [YOLOX-Darknet53](../../../exps/default/yolov3.py) | 63.72M | 185.3 | 640x640 |48.0 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_dark_openvino.tar.gz) |
15 | | [YOLOX-X](../../../exps/default/yolox_x.py) | 99.1M | 281.9 | 640x640 |51.5 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_x_openvino.tar.gz) |
16 |
17 | ## Install OpenVINO Toolkit
18 |
19 | Please visit [Openvino Homepage](https://docs.openvinotoolkit.org/latest/get_started_guides.html) for more details.
20 |
21 | ## Set up the Environment
22 |
23 | ### For Linux
24 |
25 | **Option1. Set up the environment tempororally. You need to run this command everytime you start a new shell window.**
26 |
27 | ```shell
28 | source /opt/intel/openvino_2021/bin/setupvars.sh
29 | ```
30 |
31 | **Option2. Set up the environment permenantly.**
32 |
33 | *Step1.* For Linux:
34 | ```shell
35 | vim ~/.bashrc
36 | ```
37 |
38 | *Step2.* Add the following line into your file:
39 |
40 | ```shell
41 | source /opt/intel/openvino_2021/bin/setupvars.sh
42 | ```
43 |
44 | *Step3.* Save and exit the file, then run:
45 |
46 | ```shell
47 | source ~/.bashrc
48 | ```
49 |
50 |
51 | ## Convert model
52 |
53 | 1. Export ONNX model
54 |
55 | Please refer to the [ONNX tutorial](../../ONNXRuntime). **Note that you should set --opset to 10, otherwise your next step will fail.**
56 |
57 | 2. Convert ONNX to OpenVINO
58 |
59 | ``` shell
60 | cd /openvino_2021/deployment_tools/model_optimizer
61 | ```
62 |
63 | Install requirements for convert tool
64 |
65 | ```shell
66 | sudo ./install_prerequisites/install_prerequisites_onnx.sh
67 | ```
68 |
69 | Then convert model.
70 | ```shell
71 | python3 mo.py --input_model --input_shape [--data_type FP16]
72 | ```
73 | For example:
74 | ```shell
75 | python3 mo.py --input_model yolox_tiny.onnx --input_shape [1,3,416,416] --data_type FP16
76 | ```
77 |
78 | Make sure the input shape is consistent with [those](yolox_openvino.cpp#L24-L25) in cpp file.
79 |
80 | ## Build
81 |
82 | ### Linux
83 | ```shell
84 | source /opt/intel/openvino_2021/bin/setupvars.sh
85 | mkdir build
86 | cd build
87 | cmake ..
88 | make
89 | ```
90 |
91 | ## Demo
92 |
93 | ### c++
94 |
95 | ```shell
96 | ./yolox_openvino
97 | ```
98 |
--------------------------------------------------------------------------------
/tests/utils/test_model_utils.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
5 | import unittest
6 |
7 | import torch
8 | from torch import nn
9 |
10 | from yolox.utils import adjust_status, freeze_module
11 | from yolox.exp import get_exp
12 |
13 |
14 | class TestModelUtils(unittest.TestCase):
15 |
16 | def setUp(self):
17 | self.model: nn.Module = get_exp(exp_name="yolox-s").get_model()
18 |
19 | def test_model_state_adjust_status(self):
20 | data = torch.ones(1, 10, 10, 10)
21 | # use bn since bn changes state during train/val
22 | model = nn.BatchNorm2d(10)
23 | prev_state = model.state_dict()
24 |
25 | modes = [False, True]
26 | results = [True, False]
27 |
28 | # test under train/eval mode
29 | for mode, result in zip(modes, results):
30 | with adjust_status(model, training=mode):
31 | model(data)
32 | model_state = model.state_dict()
33 | self.assertTrue(len(model_state) == len(prev_state))
34 | self.assertEqual(
35 | result,
36 | all([torch.allclose(v, model_state[k]) for k, v in prev_state.items()])
37 | )
38 |
39 | # test recurrsive context case
40 | prev_state = model.state_dict()
41 | with adjust_status(model, training=False):
42 | with adjust_status(model, training=False):
43 | model(data)
44 | model_state = model.state_dict()
45 | self.assertTrue(len(model_state) == len(prev_state))
46 | self.assertTrue(
47 | all([torch.allclose(v, model_state[k]) for k, v in prev_state.items()])
48 | )
49 |
50 | def test_model_effect_adjust_status(self):
51 | # test context effect
52 | self.model.train()
53 | with adjust_status(self.model, training=False):
54 | for module in self.model.modules():
55 | self.assertFalse(module.training)
56 | # all training after exit
57 | for module in self.model.modules():
58 | self.assertTrue(module.training)
59 |
60 | # only backbone set to eval
61 | self.model.backbone.eval()
62 | with adjust_status(self.model, training=False):
63 | for module in self.model.modules():
64 | self.assertFalse(module.training)
65 |
66 | for name, module in self.model.named_modules():
67 | if "backbone" in name:
68 | self.assertFalse(module.training)
69 | else:
70 | self.assertTrue(module.training)
71 |
72 | def test_freeze_module(self):
73 | model = nn.Sequential(
74 | nn.Conv2d(3, 10, 1),
75 | nn.BatchNorm2d(10),
76 | nn.ReLU(),
77 | )
78 | data = torch.rand(1, 3, 10, 10)
79 | model.train()
80 | assert isinstance(model[1], nn.BatchNorm2d)
81 | before_states = model[1].state_dict()
82 | freeze_module(model[1])
83 | model(data)
84 | after_states = model[1].state_dict()
85 | self.assertTrue(
86 | all([torch.allclose(v, after_states[k]) for k, v in before_states.items()])
87 | )
88 |
89 | # yolox test
90 | self.model.train()
91 | for module in self.model.modules():
92 | self.assertTrue(module.training)
93 |
94 | freeze_module(self.model, "backbone")
95 | for module in self.model.backbone.modules():
96 | self.assertFalse(module.training)
97 | for p in self.model.backbone.parameters():
98 | self.assertFalse(p.requires_grad)
99 |
100 | for module in self.model.head.modules():
101 | self.assertTrue(module.training)
102 | for p in self.model.head.parameters():
103 | self.assertTrue(p.requires_grad)
104 |
105 |
106 | if __name__ == "__main__":
107 | unittest.main()
108 |
--------------------------------------------------------------------------------
/demo/MegEngine/python/models/yolo_pafpn.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- encoding: utf-8 -*-
3 | # Copyright (c) Megvii Inc. All rights reserved.
4 |
5 | import megengine.module as M
6 | import megengine.functional as F
7 |
8 | from .darknet import CSPDarknet
9 | from .network_blocks import BaseConv, CSPLayer, DWConv, UpSample
10 |
11 |
12 | class YOLOPAFPN(M.Module):
13 | """
14 | YOLOv3 model. Darknet 53 is the default backbone of this model.
15 | """
16 |
17 | def __init__(
18 | self, depth=1.0, width=1.0, in_features=("dark3", "dark4", "dark5"),
19 | in_channels=[256, 512, 1024], depthwise=False, act="silu",
20 | ):
21 | super().__init__()
22 | self.backbone = CSPDarknet(depth, width, depthwise=depthwise, act=act)
23 | self.in_features = in_features
24 | self.in_channels = in_channels
25 | Conv = DWConv if depthwise else BaseConv
26 |
27 | self.upsample = UpSample(scale_factor=2, mode="bilinear")
28 | self.lateral_conv0 = BaseConv(
29 | int(in_channels[2] * width), int(in_channels[1] * width), 1, 1, act=act
30 | )
31 | self.C3_p4 = CSPLayer(
32 | int(2 * in_channels[1] * width),
33 | int(in_channels[1] * width),
34 | round(3 * depth),
35 | False,
36 | depthwise=depthwise,
37 | act=act,
38 | ) # cat
39 |
40 | self.reduce_conv1 = BaseConv(
41 | int(in_channels[1] * width), int(in_channels[0] * width), 1, 1, act=act
42 | )
43 | self.C3_p3 = CSPLayer(
44 | int(2 * in_channels[0] * width),
45 | int(in_channels[0] * width),
46 | round(3 * depth),
47 | False,
48 | depthwise=depthwise,
49 | act=act,
50 | )
51 |
52 | # bottom-up conv
53 | self.bu_conv2 = Conv(
54 | int(in_channels[0] * width), int(in_channels[0] * width), 3, 2, act=act
55 | )
56 | self.C3_n3 = CSPLayer(
57 | int(2 * in_channels[0] * width),
58 | int(in_channels[1] * width),
59 | round(3 * depth),
60 | False,
61 | depthwise=depthwise,
62 | act=act,
63 | )
64 |
65 | # bottom-up conv
66 | self.bu_conv1 = Conv(
67 | int(in_channels[1] * width), int(in_channels[1] * width), 3, 2, act=act
68 | )
69 | self.C3_n4 = CSPLayer(
70 | int(2 * in_channels[1] * width),
71 | int(in_channels[2] * width),
72 | round(3 * depth),
73 | False,
74 | depthwise=depthwise,
75 | act=act,
76 | )
77 |
78 | def forward(self, input):
79 | """
80 | Args:
81 | inputs: input images.
82 |
83 | Returns:
84 | Tuple[Tensor]: FPN feature.
85 | """
86 |
87 | # backbone
88 | out_features = self.backbone(input)
89 | features = [out_features[f] for f in self.in_features]
90 | [x2, x1, x0] = features
91 |
92 | fpn_out0 = self.lateral_conv0(x0) # 1024->512/32
93 | f_out0 = self.upsample(fpn_out0) # 512/16
94 | f_out0 = F.concat([f_out0, x1], 1) # 512->1024/16
95 | f_out0 = self.C3_p4(f_out0) # 1024->512/16
96 |
97 | fpn_out1 = self.reduce_conv1(f_out0) # 512->256/16
98 | f_out1 = self.upsample(fpn_out1) # 256/8
99 | f_out1 = F.concat([f_out1, x2], 1) # 256->512/8
100 | pan_out2 = self.C3_p3(f_out1) # 512->256/8
101 |
102 | p_out1 = self.bu_conv2(pan_out2) # 256->256/16
103 | p_out1 = F.concat([p_out1, fpn_out1], 1) # 256->512/16
104 | pan_out1 = self.C3_n3(p_out1) # 512->512/16
105 |
106 | p_out0 = self.bu_conv1(pan_out1) # 512->512/32
107 | p_out0 = F.concat([p_out0, fpn_out0], 1) # 512->1024/32
108 | pan_out0 = self.C3_n4(p_out0) # 1024->1024/32
109 |
110 | outputs = (pan_out2, pan_out1, pan_out0)
111 | return outputs
112 |
--------------------------------------------------------------------------------
/yolox/models/yolo_pafpn.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python
2 | # -*- encoding: utf-8 -*-
3 | # Copyright (c) Megvii Inc. All rights reserved.
4 |
5 | import torch
6 | import torch.nn as nn
7 |
8 | from .darknet import CSPDarknet
9 | from .network_blocks import BaseConv, CSPLayer, DWConv
10 |
11 |
12 | class YOLOPAFPN(nn.Module):
13 | """
14 | YOLOv3 model. Darknet 53 is the default backbone of this model.
15 | """
16 |
17 | def __init__(
18 | self,
19 | depth=1.0,
20 | width=1.0,
21 | in_features=("dark3", "dark4", "dark5"),
22 | in_channels=[256, 512, 1024],
23 | depthwise=False,
24 | act="silu",
25 | ):
26 | super().__init__()
27 | self.backbone = CSPDarknet(depth, width, depthwise=depthwise, act=act)
28 | self.in_features = in_features
29 | self.in_channels = in_channels
30 | Conv = DWConv if depthwise else BaseConv
31 |
32 | self.upsample = nn.Upsample(scale_factor=2, mode="nearest")
33 | self.lateral_conv0 = BaseConv(
34 | int(in_channels[2] * width), int(in_channels[1] * width), 1, 1, act=act
35 | )
36 | self.C3_p4 = CSPLayer(
37 | int(2 * in_channels[1] * width),
38 | int(in_channels[1] * width),
39 | round(3 * depth),
40 | False,
41 | depthwise=depthwise,
42 | act=act,
43 | ) # cat
44 |
45 | self.reduce_conv1 = BaseConv(
46 | int(in_channels[1] * width), int(in_channels[0] * width), 1, 1, act=act
47 | )
48 | self.C3_p3 = CSPLayer(
49 | int(2 * in_channels[0] * width),
50 | int(in_channels[0] * width),
51 | round(3 * depth),
52 | False,
53 | depthwise=depthwise,
54 | act=act,
55 | )
56 |
57 | # bottom-up conv
58 | self.bu_conv2 = Conv(
59 | int(in_channels[0] * width), int(in_channels[0] * width), 3, 2, act=act
60 | )
61 | self.C3_n3 = CSPLayer(
62 | int(2 * in_channels[0] * width),
63 | int(in_channels[1] * width),
64 | round(3 * depth),
65 | False,
66 | depthwise=depthwise,
67 | act=act,
68 | )
69 |
70 | # bottom-up conv
71 | self.bu_conv1 = Conv(
72 | int(in_channels[1] * width), int(in_channels[1] * width), 3, 2, act=act
73 | )
74 | self.C3_n4 = CSPLayer(
75 | int(2 * in_channels[1] * width),
76 | int(in_channels[2] * width),
77 | round(3 * depth),
78 | False,
79 | depthwise=depthwise,
80 | act=act,
81 | )
82 |
83 | def forward(self, input):
84 | """
85 | Args:
86 | inputs: input images.
87 |
88 | Returns:
89 | Tuple[Tensor]: FPN feature.
90 | """
91 |
92 | # backbone
93 | out_features = self.backbone(input)
94 | features = [out_features[f] for f in self.in_features]
95 | [x2, x1, x0] = features
96 |
97 | fpn_out0 = self.lateral_conv0(x0) # 1024->512/32
98 | f_out0 = self.upsample(fpn_out0) # 512/16
99 | f_out0 = torch.cat([f_out0, x1], 1) # 512->1024/16
100 | f_out0 = self.C3_p4(f_out0) # 1024->512/16
101 |
102 | fpn_out1 = self.reduce_conv1(f_out0) # 512->256/16
103 | f_out1 = self.upsample(fpn_out1) # 256/8
104 | f_out1 = torch.cat([f_out1, x2], 1) # 256->512/8
105 | pan_out2 = self.C3_p3(f_out1) # 512->256/8
106 |
107 | p_out1 = self.bu_conv2(pan_out2) # 256->256/16
108 | p_out1 = torch.cat([p_out1, fpn_out1], 1) # 256->512/16
109 | pan_out1 = self.C3_n3(p_out1) # 512->512/16
110 |
111 | p_out0 = self.bu_conv1(pan_out1) # 512->512/32
112 | p_out0 = torch.cat([p_out0, fpn_out0], 1) # 512->1024/32
113 | pan_out0 = self.C3_n4(p_out0) # 1024->1024/32
114 |
115 | outputs = (pan_out2, pan_out1, pan_out0)
116 | return outputs
117 |
--------------------------------------------------------------------------------
/yolox/utils/metric.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding: utf-8 -*-
3 | # Copyright (c) Megvii Inc. All rights reserved.
4 | import functools
5 | import os
6 | import time
7 | from collections import defaultdict, deque
8 | import psutil
9 |
10 | import numpy as np
11 |
12 | import torch
13 |
14 | __all__ = [
15 | "AverageMeter",
16 | "MeterBuffer",
17 | "get_total_and_free_memory_in_Mb",
18 | "occupy_mem",
19 | "gpu_mem_usage",
20 | "mem_usage"
21 | ]
22 |
23 |
24 | def get_total_and_free_memory_in_Mb(cuda_device):
25 | devices_info_str = os.popen(
26 | "nvidia-smi --query-gpu=memory.total,memory.used --format=csv,nounits,noheader"
27 | )
28 | devices_info = devices_info_str.read().strip().split("\n")
29 | if "CUDA_VISIBLE_DEVICES" in os.environ:
30 | visible_devices = os.environ["CUDA_VISIBLE_DEVICES"].split(',')
31 | cuda_device = int(visible_devices[cuda_device])
32 | total, used = devices_info[int(cuda_device)].split(",")
33 | return int(total), int(used)
34 |
35 |
36 | def occupy_mem(cuda_device, mem_ratio=0.9):
37 | """
38 | pre-allocate gpu memory for training to avoid memory Fragmentation.
39 | """
40 | total, used = get_total_and_free_memory_in_Mb(cuda_device)
41 | max_mem = int(total * mem_ratio)
42 | block_mem = max_mem - used
43 | x = torch.cuda.FloatTensor(256, 1024, block_mem)
44 | del x
45 | time.sleep(5)
46 |
47 |
48 | def gpu_mem_usage():
49 | """
50 | Compute the GPU memory usage for the current device (MB).
51 | """
52 | mem_usage_bytes = torch.cuda.max_memory_allocated()
53 | return mem_usage_bytes / (1024 * 1024)
54 |
55 |
56 | def mem_usage():
57 | """
58 | Compute the memory usage for the current machine (GB).
59 | """
60 | gb = 1 << 30
61 | mem = psutil.virtual_memory()
62 | return mem.used / gb
63 |
64 |
65 | class AverageMeter:
66 | """Track a series of values and provide access to smoothed values over a
67 | window or the global series average.
68 | """
69 |
70 | def __init__(self, window_size=50):
71 | self._deque = deque(maxlen=window_size)
72 | self._total = 0.0
73 | self._count = 0
74 |
75 | def update(self, value):
76 | self._deque.append(value)
77 | self._count += 1
78 | self._total += value
79 |
80 | @property
81 | def median(self):
82 | d = np.array(list(self._deque))
83 | return np.median(d)
84 |
85 | @property
86 | def avg(self):
87 | # if deque is empty, nan will be returned.
88 | d = np.array(list(self._deque))
89 | return d.mean()
90 |
91 | @property
92 | def global_avg(self):
93 | return self._total / max(self._count, 1e-5)
94 |
95 | @property
96 | def latest(self):
97 | return self._deque[-1] if len(self._deque) > 0 else None
98 |
99 | @property
100 | def total(self):
101 | return self._total
102 |
103 | def reset(self):
104 | self._deque.clear()
105 | self._total = 0.0
106 | self._count = 0
107 |
108 | def clear(self):
109 | self._deque.clear()
110 |
111 |
112 | class MeterBuffer(defaultdict):
113 | """Computes and stores the average and current value"""
114 |
115 | def __init__(self, window_size=20):
116 | factory = functools.partial(AverageMeter, window_size=window_size)
117 | super().__init__(factory)
118 |
119 | def reset(self):
120 | for v in self.values():
121 | v.reset()
122 |
123 | def get_filtered_meter(self, filter_key="time"):
124 | return {k: v for k, v in self.items() if filter_key in k}
125 |
126 | def update(self, values=None, **kwargs):
127 | if values is None:
128 | values = {}
129 | values.update(kwargs)
130 | for k, v in values.items():
131 | if isinstance(v, torch.Tensor):
132 | v = v.detach()
133 | self[k].update(v)
134 |
135 | def clear_meters(self):
136 | for v in self.values():
137 | v.clear()
138 |
--------------------------------------------------------------------------------
/tools/export_onnx.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
5 | import argparse
6 | import os
7 | from loguru import logger
8 |
9 | import torch
10 | from torch import nn
11 |
12 | from yolox.exp import get_exp
13 | from yolox.models.network_blocks import SiLU
14 | from yolox.utils import replace_module
15 |
16 |
17 | def make_parser():
18 | parser = argparse.ArgumentParser("YOLOX onnx deploy")
19 | parser.add_argument(
20 | "--output-name", type=str, default="yolox.onnx", help="output name of models"
21 | )
22 | parser.add_argument(
23 | "--input", default="images", type=str, help="input node name of onnx model"
24 | )
25 | parser.add_argument(
26 | "--output", default="output", type=str, help="output node name of onnx model"
27 | )
28 | parser.add_argument(
29 | "-o", "--opset", default=11, type=int, help="onnx opset version"
30 | )
31 | parser.add_argument("--batch-size", type=int, default=1, help="batch size")
32 | parser.add_argument(
33 | "--dynamic", action="store_true", help="whether the input shape should be dynamic or not"
34 | )
35 | parser.add_argument("--no-onnxsim", action="store_true", help="use onnxsim or not")
36 | parser.add_argument(
37 | "-f",
38 | "--exp_file",
39 | default=None,
40 | type=str,
41 | help="experiment description file",
42 | )
43 | parser.add_argument("-expn", "--experiment-name", type=str, default=None)
44 | parser.add_argument("-n", "--name", type=str, default=None, help="model name")
45 | parser.add_argument("-c", "--ckpt", default=None, type=str, help="ckpt path")
46 | parser.add_argument(
47 | "opts",
48 | help="Modify config options using the command-line",
49 | default=None,
50 | nargs=argparse.REMAINDER,
51 | )
52 | parser.add_argument(
53 | "--decode_in_inference",
54 | action="store_true",
55 | help="decode in inference or not"
56 | )
57 |
58 | return parser
59 |
60 |
61 | @logger.catch
62 | def main():
63 | args = make_parser().parse_args()
64 | logger.info("args value: {}".format(args))
65 | exp = get_exp(args.exp_file, args.name)
66 | exp.merge(args.opts)
67 |
68 | if not args.experiment_name:
69 | args.experiment_name = exp.exp_name
70 |
71 | model = exp.get_model()
72 | if args.ckpt is None:
73 | file_name = os.path.join(exp.output_dir, args.experiment_name)
74 | ckpt_file = os.path.join(file_name, "best_ckpt.pth")
75 | else:
76 | ckpt_file = args.ckpt
77 |
78 | # load the model state dict
79 | ckpt = torch.load(ckpt_file, map_location="cpu")
80 |
81 | model.eval()
82 | if "model" in ckpt:
83 | ckpt = ckpt["model"]
84 | model.load_state_dict(ckpt)
85 | model = replace_module(model, nn.SiLU, SiLU)
86 | model.head.decode_in_inference = args.decode_in_inference
87 |
88 | logger.info("loading checkpoint done.")
89 | dummy_input = torch.randn(args.batch_size, 3, exp.test_size[0], exp.test_size[1])
90 |
91 | torch.onnx._export(
92 | model,
93 | dummy_input,
94 | args.output_name,
95 | input_names=[args.input],
96 | output_names=[args.output],
97 | dynamic_axes={args.input: {0: 'batch'},
98 | args.output: {0: 'batch'}} if args.dynamic else None,
99 | opset_version=args.opset,
100 | )
101 | logger.info("generated onnx model named {}".format(args.output_name))
102 |
103 | if not args.no_onnxsim:
104 | import onnx
105 | from onnxsim import simplify
106 |
107 | # use onnx-simplifier to reduce reduent model.
108 | onnx_model = onnx.load(args.output_name)
109 | model_simp, check = simplify(onnx_model)
110 | assert check, "Simplified ONNX model could not be validated"
111 | onnx.save(model_simp, args.output_name)
112 | logger.info("generated simplified onnx model named {}".format(args.output_name))
113 |
114 |
115 | if __name__ == "__main__":
116 | main()
117 |
--------------------------------------------------------------------------------
/yolox/data/dataloading.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
5 | import os
6 | import random
7 | import uuid
8 |
9 | import numpy as np
10 |
11 | import torch
12 | from torch.utils.data.dataloader import DataLoader as torchDataLoader
13 | from torch.utils.data.dataloader import default_collate
14 |
15 | from .samplers import YoloBatchSampler
16 |
17 |
18 | def get_yolox_datadir():
19 | """
20 | get dataset dir of YOLOX. If environment variable named `YOLOX_DATADIR` is set,
21 | this function will return value of the environment variable. Otherwise, use data
22 | """
23 | yolox_datadir = os.getenv("YOLOX_DATADIR", None)
24 | if yolox_datadir is None:
25 | import yolox
26 |
27 | yolox_path = os.path.dirname(os.path.dirname(yolox.__file__))
28 | yolox_datadir = os.path.join(yolox_path, "datasets")
29 | return yolox_datadir
30 |
31 |
32 | class DataLoader(torchDataLoader):
33 | """
34 | Lightnet dataloader that enables on the fly resizing of the images.
35 | See :class:`torch.utils.data.DataLoader` for more information on the arguments.
36 | Check more on the following website:
37 | https://gitlab.com/EAVISE/lightnet/-/blob/master/lightnet/data/_dataloading.py
38 | """
39 |
40 | def __init__(self, *args, **kwargs):
41 | super().__init__(*args, **kwargs)
42 | self.__initialized = False
43 | shuffle = False
44 | batch_sampler = None
45 | if len(args) > 5:
46 | shuffle = args[2]
47 | sampler = args[3]
48 | batch_sampler = args[4]
49 | elif len(args) > 4:
50 | shuffle = args[2]
51 | sampler = args[3]
52 | if "batch_sampler" in kwargs:
53 | batch_sampler = kwargs["batch_sampler"]
54 | elif len(args) > 3:
55 | shuffle = args[2]
56 | if "sampler" in kwargs:
57 | sampler = kwargs["sampler"]
58 | if "batch_sampler" in kwargs:
59 | batch_sampler = kwargs["batch_sampler"]
60 | else:
61 | if "shuffle" in kwargs:
62 | shuffle = kwargs["shuffle"]
63 | if "sampler" in kwargs:
64 | sampler = kwargs["sampler"]
65 | if "batch_sampler" in kwargs:
66 | batch_sampler = kwargs["batch_sampler"]
67 |
68 | # Use custom BatchSampler
69 | if batch_sampler is None:
70 | if sampler is None:
71 | if shuffle:
72 | sampler = torch.utils.data.sampler.RandomSampler(self.dataset)
73 | # sampler = torch.utils.data.DistributedSampler(self.dataset)
74 | else:
75 | sampler = torch.utils.data.sampler.SequentialSampler(self.dataset)
76 | batch_sampler = YoloBatchSampler(
77 | sampler,
78 | self.batch_size,
79 | self.drop_last,
80 | input_dimension=self.dataset.input_dim,
81 | )
82 | # batch_sampler = IterationBasedBatchSampler(batch_sampler, num_iterations =
83 |
84 | self.batch_sampler = batch_sampler
85 |
86 | self.__initialized = True
87 |
88 | def close_mosaic(self):
89 | self.batch_sampler.mosaic = False
90 |
91 |
92 | def list_collate(batch):
93 | """
94 | Function that collates lists or tuples together into one list (of lists/tuples).
95 | Use this as the collate function in a Dataloader, if you want to have a list of
96 | items as an output, as opposed to tensors (eg. Brambox.boxes).
97 | """
98 | items = list(zip(*batch))
99 |
100 | for i in range(len(items)):
101 | if isinstance(items[i][0], (list, tuple)):
102 | items[i] = list(items[i])
103 | else:
104 | items[i] = default_collate(items[i])
105 |
106 | return items
107 |
108 |
109 | def worker_init_reset_seed(worker_id):
110 | seed = uuid.uuid4().int % 2**32
111 | random.seed(seed)
112 | torch.set_rng_state(torch.manual_seed(seed).get_state())
113 | np.random.seed(seed)
114 |
--------------------------------------------------------------------------------
/yolox/utils/visualize.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii Inc. All rights reserved.
4 |
5 | import cv2
6 | import numpy as np
7 |
8 | __all__ = ["vis"]
9 |
10 |
11 | def vis(img, boxes, scores, cls_ids, conf=0.5, class_names=None):
12 |
13 | for i in range(len(boxes)):
14 | box = boxes[i]
15 | cls_id = int(cls_ids[i])
16 | score = scores[i]
17 | if score < conf:
18 | continue
19 | x0 = int(box[0])
20 | y0 = int(box[1])
21 | x1 = int(box[2])
22 | y1 = int(box[3])
23 |
24 | color = (_COLORS[cls_id] * 255).astype(np.uint8).tolist()
25 | text = '{}:{:.1f}%'.format(class_names[cls_id], score * 100)
26 | txt_color = (0, 0, 0) if np.mean(_COLORS[cls_id]) > 0.5 else (255, 255, 255)
27 | font = cv2.FONT_HERSHEY_SIMPLEX
28 |
29 | txt_size = cv2.getTextSize(text, font, 0.4, 1)[0]
30 | cv2.rectangle(img, (x0, y0), (x1, y1), color, 2)
31 |
32 | txt_bk_color = (_COLORS[cls_id] * 255 * 0.7).astype(np.uint8).tolist()
33 | cv2.rectangle(
34 | img,
35 | (x0, y0 + 1),
36 | (x0 + txt_size[0] + 1, y0 + int(1.5*txt_size[1])),
37 | txt_bk_color,
38 | -1
39 | )
40 | cv2.putText(img, text, (x0, y0 + txt_size[1]), font, 0.4, txt_color, thickness=1)
41 |
42 | return img
43 |
44 |
45 | _COLORS = np.array(
46 | [
47 | 0.000, 0.447, 0.741,
48 | 0.850, 0.325, 0.098,
49 | 0.929, 0.694, 0.125,
50 | 0.494, 0.184, 0.556,
51 | 0.466, 0.674, 0.188,
52 | 0.301, 0.745, 0.933,
53 | 0.635, 0.078, 0.184,
54 | 0.300, 0.300, 0.300,
55 | 0.600, 0.600, 0.600,
56 | 1.000, 0.000, 0.000,
57 | 1.000, 0.500, 0.000,
58 | 0.749, 0.749, 0.000,
59 | 0.000, 1.000, 0.000,
60 | 0.000, 0.000, 1.000,
61 | 0.667, 0.000, 1.000,
62 | 0.333, 0.333, 0.000,
63 | 0.333, 0.667, 0.000,
64 | 0.333, 1.000, 0.000,
65 | 0.667, 0.333, 0.000,
66 | 0.667, 0.667, 0.000,
67 | 0.667, 1.000, 0.000,
68 | 1.000, 0.333, 0.000,
69 | 1.000, 0.667, 0.000,
70 | 1.000, 1.000, 0.000,
71 | 0.000, 0.333, 0.500,
72 | 0.000, 0.667, 0.500,
73 | 0.000, 1.000, 0.500,
74 | 0.333, 0.000, 0.500,
75 | 0.333, 0.333, 0.500,
76 | 0.333, 0.667, 0.500,
77 | 0.333, 1.000, 0.500,
78 | 0.667, 0.000, 0.500,
79 | 0.667, 0.333, 0.500,
80 | 0.667, 0.667, 0.500,
81 | 0.667, 1.000, 0.500,
82 | 1.000, 0.000, 0.500,
83 | 1.000, 0.333, 0.500,
84 | 1.000, 0.667, 0.500,
85 | 1.000, 1.000, 0.500,
86 | 0.000, 0.333, 1.000,
87 | 0.000, 0.667, 1.000,
88 | 0.000, 1.000, 1.000,
89 | 0.333, 0.000, 1.000,
90 | 0.333, 0.333, 1.000,
91 | 0.333, 0.667, 1.000,
92 | 0.333, 1.000, 1.000,
93 | 0.667, 0.000, 1.000,
94 | 0.667, 0.333, 1.000,
95 | 0.667, 0.667, 1.000,
96 | 0.667, 1.000, 1.000,
97 | 1.000, 0.000, 1.000,
98 | 1.000, 0.333, 1.000,
99 | 1.000, 0.667, 1.000,
100 | 0.333, 0.000, 0.000,
101 | 0.500, 0.000, 0.000,
102 | 0.667, 0.000, 0.000,
103 | 0.833, 0.000, 0.000,
104 | 1.000, 0.000, 0.000,
105 | 0.000, 0.167, 0.000,
106 | 0.000, 0.333, 0.000,
107 | 0.000, 0.500, 0.000,
108 | 0.000, 0.667, 0.000,
109 | 0.000, 0.833, 0.000,
110 | 0.000, 1.000, 0.000,
111 | 0.000, 0.000, 0.167,
112 | 0.000, 0.000, 0.333,
113 | 0.000, 0.000, 0.500,
114 | 0.000, 0.000, 0.667,
115 | 0.000, 0.000, 0.833,
116 | 0.000, 0.000, 1.000,
117 | 0.000, 0.000, 0.000,
118 | 0.143, 0.143, 0.143,
119 | 0.286, 0.286, 0.286,
120 | 0.429, 0.429, 0.429,
121 | 0.571, 0.571, 0.571,
122 | 0.714, 0.714, 0.714,
123 | 0.857, 0.857, 0.857,
124 | 0.000, 0.447, 0.741,
125 | 0.314, 0.717, 0.741,
126 | 0.50, 0.5, 0
127 | ]
128 | ).astype(np.float32).reshape(-1, 3)
129 |
--------------------------------------------------------------------------------
/yolox/layers/cocoeval/cocoeval.h:
--------------------------------------------------------------------------------
1 | // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
2 | #pragma once
3 |
4 | #include
5 | #include
6 | #include
7 | #include
8 | #include
9 |
10 | namespace py = pybind11;
11 |
12 | namespace COCOeval {
13 |
14 | // Annotation data for a single object instance in an image
15 | struct InstanceAnnotation {
16 | InstanceAnnotation(
17 | uint64_t id,
18 | double score,
19 | double area,
20 | bool is_crowd,
21 | bool ignore)
22 | : id{id}, score{score}, area{area}, is_crowd{is_crowd}, ignore{ignore} {}
23 | uint64_t id;
24 | double score = 0.;
25 | double area = 0.;
26 | bool is_crowd = false;
27 | bool ignore = false;
28 | };
29 |
30 | // Stores intermediate results for evaluating detection results for a single
31 | // image that has D detected instances and G ground truth instances. This stores
32 | // matches between detected and ground truth instances
33 | struct ImageEvaluation {
34 | // For each of the D detected instances, the id of the matched ground truth
35 | // instance, or 0 if unmatched
36 | std::vector detection_matches;
37 |
38 | // The detection score of each of the D detected instances
39 | std::vector detection_scores;
40 |
41 | // Marks whether or not each of G instances was ignored from evaluation (e.g.,
42 | // because it's outside area_range)
43 | std::vector ground_truth_ignores;
44 |
45 | // Marks whether or not each of D instances was ignored from evaluation (e.g.,
46 | // because it's outside aRng)
47 | std::vector detection_ignores;
48 | };
49 |
50 | template
51 | using ImageCategoryInstances = std::vector>>;
52 |
53 | // C++ implementation of COCO API cocoeval.py::COCOeval.evaluateImg(). For each
54 | // combination of image, category, area range settings, and IOU thresholds to
55 | // evaluate, it matches detected instances to ground truth instances and stores
56 | // the results into a vector of ImageEvaluation results, which will be
57 | // interpreted by the COCOeval::Accumulate() function to produce precion-recall
58 | // curves. The parameters of nested vectors have the following semantics:
59 | // image_category_ious[i][c][d][g] is the intersection over union of the d'th
60 | // detected instance and g'th ground truth instance of
61 | // category category_ids[c] in image image_ids[i]
62 | // image_category_ground_truth_instances[i][c] is a vector of ground truth
63 | // instances in image image_ids[i] of category category_ids[c]
64 | // image_category_detection_instances[i][c] is a vector of detected
65 | // instances in image image_ids[i] of category category_ids[c]
66 | std::vector EvaluateImages(
67 | const std::vector>& area_ranges, // vector of 2-tuples
68 | int max_detections,
69 | const std::vector& iou_thresholds,
70 | const ImageCategoryInstances>& image_category_ious,
71 | const ImageCategoryInstances&
72 | image_category_ground_truth_instances,
73 | const ImageCategoryInstances&
74 | image_category_detection_instances);
75 |
76 | // C++ implementation of COCOeval.accumulate(), which generates precision
77 | // recall curves for each set of category, IOU threshold, detection area range,
78 | // and max number of detections parameters. It is assumed that the parameter
79 | // evaluations is the return value of the functon COCOeval::EvaluateImages(),
80 | // which was called with the same parameter settings params
81 | py::dict Accumulate(
82 | const py::object& params,
83 | const std::vector& evalutations);
84 |
85 | } // namespace COCOeval
86 |
87 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
88 | {
89 | m.def("COCOevalAccumulate", &COCOeval::Accumulate, "COCOeval::Accumulate");
90 | m.def(
91 | "COCOevalEvaluateImages",
92 | &COCOeval::EvaluateImages,
93 | "COCOeval::EvaluateImages");
94 | pybind11::class_(m, "InstanceAnnotation")
95 | .def(pybind11::init());
96 | pybind11::class_(m, "ImageEvaluation")
97 | .def(pybind11::init<>());
98 | }
99 |
--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
1 | ### Linux ###
2 | *~
3 |
4 | # user experiments directory
5 | YOLOX_outputs/
6 | datasets/
7 | # do not ignore datasets under yolox/data
8 | !*yolox/data/datasets/
9 |
10 | # temporary files which can be created if a process still has a handle open of a deleted file
11 | .fuse_hidden*
12 |
13 | # KDE directory preferences
14 | .directory
15 |
16 | # Linux trash folder which might appear on any partition or disk
17 | .Trash-*
18 |
19 | # .nfs files are created when an open file is removed but is still being accessed
20 | .nfs*
21 |
22 | ### PyCharm ###
23 | # User-specific stuff
24 | .idea
25 |
26 | # CMake
27 | cmake-build-*/
28 |
29 | # Mongo Explorer plugin
30 | .idea/**/mongoSettings.xml
31 |
32 | # File-based project format
33 | *.iws
34 |
35 | # IntelliJ
36 | out/
37 |
38 | # mpeltonen/sbt-idea plugin
39 | .idea_modules/
40 |
41 | # JIRA plugin
42 | atlassian-ide-plugin.xml
43 |
44 | # Cursive Clojure plugin
45 | .idea/replstate.xml
46 |
47 | # Crashlytics plugin (for Android Studio and IntelliJ)
48 | com_crashlytics_export_strings.xml
49 | crashlytics.properties
50 | crashlytics-build.properties
51 | fabric.properties
52 |
53 | # Editor-based Rest Client
54 | .idea/httpRequests
55 |
56 | # Android studio 3.1+ serialized cache file
57 | .idea/caches/build_file_checksums.ser
58 |
59 | # JetBrains templates
60 | **___jb_tmp___
61 |
62 | ### Python ###
63 | # Byte-compiled / optimized / DLL files
64 | __pycache__/
65 | *.py[cod]
66 | *$py.class
67 |
68 | # C extensions
69 | *.so
70 |
71 | # Distribution / packaging
72 | .Python
73 | build/
74 | develop-eggs/
75 | dist/
76 | downloads/
77 | eggs/
78 | .eggs/
79 | lib/
80 | lib64/
81 | parts/
82 | sdist/
83 | var/
84 | wheels/
85 | pip-wheel-metadata/
86 | share/python-wheels/
87 | *.egg-info/
88 | .installed.cfg
89 | *.egg
90 | MANIFEST
91 |
92 | # PyInstaller
93 | # Usually these files are written by a python script from a template
94 | # before PyInstaller builds the exe, so as to inject date/other infos into it.
95 | *.manifest
96 | *.spec
97 |
98 | # Installer logs
99 | pip-log.txt
100 | pip-delete-this-directory.txt
101 |
102 | # Unit test / coverage reports
103 | htmlcov/
104 | .tox/
105 | .nox/
106 | .coverage
107 | .coverage.*
108 | .cache
109 | nosetests.xml
110 | coverage.xml
111 | *.cover
112 | .hypothesis/
113 | .pytest_cache/
114 |
115 | # Translations
116 | *.mo
117 | *.pot
118 |
119 | # Django stuff:
120 | *.log
121 | local_settings.py
122 | db.sqlite3
123 |
124 | # Flask stuff:
125 | instance/
126 | .webassets-cache
127 |
128 | # Scrapy stuff:
129 | .scrapy
130 |
131 | # Sphinx documentation
132 | docs/_build/
133 | docs/build/
134 |
135 | # PyBuilder
136 | target/
137 |
138 | # Jupyter Notebook
139 | .ipynb_checkpoints
140 |
141 | # IPython
142 | profile_default/
143 | ipython_config.py
144 |
145 | # pyenv
146 | .python-version
147 |
148 | # pipenv
149 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
150 | # However, in case of collaboration, if having platform-specific dependencies or dependencies
151 | # having no cross-platform support, pipenv may install dependencies that don’t work, or not
152 | # install all needed dependencies.
153 | #Pipfile.lock
154 |
155 | # celery beat schedule file
156 | celerybeat-schedule
157 |
158 | # SageMath parsed files
159 | *.sage.py
160 |
161 | # Environments
162 | .env
163 | .venv
164 | env/
165 | venv/
166 | ENV/
167 | env.bak/
168 | venv.bak/
169 |
170 | # Spyder project settings
171 | .spyderproject
172 | .spyproject
173 |
174 | # Rope project settings
175 | .ropeproject
176 |
177 | # mkdocs documentation
178 | /site
179 |
180 | # mypy
181 | .mypy_cache/
182 | .dmypy.json
183 | dmypy.json
184 |
185 | # Pyre type checker
186 | .pyre/
187 |
188 | ### Vim ###
189 | # Swap
190 | [._]*.s[a-v][a-z]
191 | [._]*.sw[a-p]
192 | [._]s[a-rt-v][a-z]
193 | [._]ss[a-gi-z]
194 | [._]sw[a-p]
195 |
196 | # Session
197 | Session.vim
198 |
199 | # Temporary
200 | .netrwhist
201 | # Auto-generated tag files
202 | tags
203 | # Persistent undo
204 | [._]*.un~
205 |
206 | # output
207 | docs/api
208 | .code-workspace.code-workspace
209 | *.pkl
210 | *.npy
211 | *.pth
212 | *.onnx
213 | *.engine
214 | events.out.tfevents*
215 |
216 | # vscode
217 | *.code-workspace
218 | .vscode
219 |
220 | # vim
221 | .vim
222 |
223 | # OS generated files
224 | .DS_Store
225 | .DS_Store?
226 | .Trashes
227 | ehthumbs.db
228 | Thumbs.db
229 |
--------------------------------------------------------------------------------
/docs/quick_run.md:
--------------------------------------------------------------------------------
1 |
2 | # Get Started
3 |
4 | ## 1.Installation
5 |
6 | Step1. Install YOLOX.
7 | ```shell
8 | git clone git@github.com:Megvii-BaseDetection/YOLOX.git
9 | cd YOLOX
10 | pip3 install -U pip && pip3 install -r requirements.txt
11 | pip3 install -v -e . # or python3 setup.py develop
12 | ```
13 | Step2. Install [pycocotools](https://github.com/cocodataset/cocoapi).
14 |
15 | ```shell
16 | pip3 install cython; pip3 install 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'
17 | ```
18 |
19 | ## 2.Demo
20 |
21 | Step1. Download a pretrained model from the benchmark table.
22 |
23 | Step2. Use either -n or -f to specify your detector's config. For example:
24 |
25 | ```shell
26 | python tools/demo.py image -n yolox-s -c /path/to/your/yolox_s.pth --path assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result --device [cpu/gpu]
27 | ```
28 | or
29 | ```shell
30 | python tools/demo.py image -f exps/default/yolox_s.py -c /path/to/your/yolox_s.pth --path assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result --device [cpu/gpu]
31 | ```
32 | Demo for video:
33 | ```shell
34 | python tools/demo.py video -n yolox-s -c /path/to/your/yolox_s.pth --path /path/to/your/video --conf 0.25 --nms 0.45 --tsize 640 --save_result --device [cpu/gpu]
35 | ```
36 |
37 |
38 | ## 3.Reproduce our results on COCO
39 |
40 | Step1. Prepare COCO dataset
41 | ```shell
42 | cd
43 | ln -s /path/to/your/COCO ./datasets/COCO
44 | ```
45 |
46 | Step2. Reproduce our results on COCO by specifying -n:
47 |
48 | ```shell
49 | python tools/train.py -n yolox-s -d 8 -b 64 --fp16 -o [--cache]
50 | yolox-m
51 | yolox-l
52 | yolox-x
53 | ```
54 | * -d: number of gpu devices
55 | * -b: total batch size, the recommended number for -b is num-gpu * 8
56 | * --fp16: mixed precision training
57 | * --cache: caching imgs into RAM to accelarate training, which need large system RAM.
58 |
59 | **Weights & Biases for Logging**
60 |
61 | To use W&B for logging, install wandb in your environment and log in to your W&B account using
62 |
63 | ```shell
64 | pip install wandb
65 | wandb login
66 | ```
67 |
68 | Log in to your W&B account
69 |
70 | To start logging metrics to W&B during training add the flag `--logger` to the previous command and use the prefix "wandb-" to specify arguments for initializing the wandb run.
71 |
72 | ```shell
73 | python tools/train.py -n yolox-s -d 8 -b 64 --fp16 -o [--cache] --logger wandb wandb-project
74 | yolox-m
75 | yolox-l
76 | yolox-x
77 | ```
78 |
79 | More WandbLogger arguments include
80 |
81 | ```shell
82 | python tools/train.py .... --logger wandb wandb-project \
83 | wandb-name \
84 | wandb-id \
85 | wandb-save_dir \
86 | wandb-num_eval_images \
87 | wandb-log_checkpoints
88 | ```
89 |
90 | More information available [here](https://docs.wandb.ai/guides/integrations/other/yolox).
91 |
92 | **Multi Machine Training**
93 |
94 | We also support multi-nodes training. Just add the following args:
95 | * --num\_machines: num of your total training nodes
96 | * --machine\_rank: specify the rank of each node
97 |
98 | When using -f, the above commands are equivalent to:
99 |
100 | ```shell
101 | python tools/train.py -f exps/default/yolox-s.py -d 8 -b 64 --fp16 -o [--cache]
102 | exps/default/yolox-m.py
103 | exps/default/yolox-l.py
104 | exps/default/yolox-x.py
105 | ```
106 |
107 | ## 4.Evaluation
108 |
109 | We support batch testing for fast evaluation:
110 |
111 | ```shell
112 | python tools/eval.py -n yolox-s -c yolox_s.pth -b 64 -d 8 --conf 0.001 [--fp16] [--fuse]
113 | yolox-m
114 | yolox-l
115 | yolox-x
116 | ```
117 | * --fuse: fuse conv and bn
118 | * -d: number of GPUs used for evaluation. DEFAULT: All GPUs available will be used.
119 | * -b: total batch size across on all GPUs
120 |
121 | To reproduce speed test, we use the following command:
122 | ```shell
123 | python tools/eval.py -n yolox-s -c yolox_s.pth -b 1 -d 1 --conf 0.001 --fp16 --fuse
124 | yolox-m
125 | yolox-l
126 | yolox-x
127 | ```
128 |
--------------------------------------------------------------------------------
/yolox/models/build.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 |
4 | import torch
5 | from torch import nn
6 | from torch.hub import load_state_dict_from_url
7 |
8 | __all__ = [
9 | "create_yolox_model",
10 | "yolox_nano",
11 | "yolox_tiny",
12 | "yolox_s",
13 | "yolox_m",
14 | "yolox_l",
15 | "yolox_x",
16 | "yolov3",
17 | "yolox_custom"
18 | ]
19 |
20 | _CKPT_ROOT_URL = "https://github.com/Megvii-BaseDetection/YOLOX/releases/download"
21 | _CKPT_FULL_PATH = {
22 | "yolox-nano": f"{_CKPT_ROOT_URL}/0.1.1rc0/yolox_nano.pth",
23 | "yolox-tiny": f"{_CKPT_ROOT_URL}/0.1.1rc0/yolox_tiny.pth",
24 | "yolox-s": f"{_CKPT_ROOT_URL}/0.1.1rc0/yolox_s.pth",
25 | "yolox-m": f"{_CKPT_ROOT_URL}/0.1.1rc0/yolox_m.pth",
26 | "yolox-l": f"{_CKPT_ROOT_URL}/0.1.1rc0/yolox_l.pth",
27 | "yolox-x": f"{_CKPT_ROOT_URL}/0.1.1rc0/yolox_x.pth",
28 | "yolov3": f"{_CKPT_ROOT_URL}/0.1.1rc0/yolox_darknet.pth",
29 | }
30 |
31 |
32 | def create_yolox_model(name: str, pretrained: bool = True, num_classes: int = 80, device=None,
33 | exp_path: str = None, ckpt_path: str = None) -> nn.Module:
34 | """creates and loads a YOLOX model
35 |
36 | Args:
37 | name (str): name of model. for example, "yolox-s", "yolox-tiny" or "yolox_custom"
38 | if you want to load your own model.
39 | pretrained (bool): load pretrained weights into the model. Default to True.
40 | device (str): default device to for model. Default to None.
41 | num_classes (int): number of model classes. Default to 80.
42 | exp_path (str): path to your own experiment file. Required if name="yolox_custom"
43 | ckpt_path (str): path to your own ckpt. Required if name="yolox_custom" and you want to
44 | load a pretrained model
45 |
46 |
47 | Returns:
48 | YOLOX model (nn.Module)
49 | """
50 | from yolox.exp import get_exp, Exp
51 |
52 | if device is None:
53 | device = "cuda:0" if torch.cuda.is_available() else "cpu"
54 | device = torch.device(device)
55 |
56 | assert name in _CKPT_FULL_PATH or name == "yolox_custom", \
57 | f"user should use one of value in {_CKPT_FULL_PATH.keys()} or \"yolox_custom\""
58 | if name in _CKPT_FULL_PATH:
59 | exp: Exp = get_exp(exp_name=name)
60 | exp.num_classes = num_classes
61 | yolox_model = exp.get_model()
62 | if pretrained and num_classes == 80:
63 | weights_url = _CKPT_FULL_PATH[name]
64 | ckpt = load_state_dict_from_url(weights_url, map_location="cpu")
65 | if "model" in ckpt:
66 | ckpt = ckpt["model"]
67 | yolox_model.load_state_dict(ckpt)
68 | else:
69 | assert exp_path is not None, "for a \"yolox_custom\" model exp_path must be provided"
70 | exp: Exp = get_exp(exp_file=exp_path)
71 | yolox_model = exp.get_model()
72 | if ckpt_path:
73 | ckpt = torch.load(ckpt_path, map_location="cpu")
74 | if "model" in ckpt:
75 | ckpt = ckpt["model"]
76 | yolox_model.load_state_dict(ckpt)
77 |
78 | yolox_model.to(device)
79 | return yolox_model
80 |
81 |
82 | def yolox_nano(pretrained: bool = True, num_classes: int = 80, device: str = None) -> nn.Module:
83 | return create_yolox_model("yolox-nano", pretrained, num_classes, device)
84 |
85 |
86 | def yolox_tiny(pretrained: bool = True, num_classes: int = 80, device: str = None) -> nn.Module:
87 | return create_yolox_model("yolox-tiny", pretrained, num_classes, device)
88 |
89 |
90 | def yolox_s(pretrained: bool = True, num_classes: int = 80, device: str = None) -> nn.Module:
91 | return create_yolox_model("yolox-s", pretrained, num_classes, device)
92 |
93 |
94 | def yolox_m(pretrained: bool = True, num_classes: int = 80, device: str = None) -> nn.Module:
95 | return create_yolox_model("yolox-m", pretrained, num_classes, device)
96 |
97 |
98 | def yolox_l(pretrained: bool = True, num_classes: int = 80, device: str = None) -> nn.Module:
99 | return create_yolox_model("yolox-l", pretrained, num_classes, device)
100 |
101 |
102 | def yolox_x(pretrained: bool = True, num_classes: int = 80, device: str = None) -> nn.Module:
103 | return create_yolox_model("yolox-x", pretrained, num_classes, device)
104 |
105 |
106 | def yolov3(pretrained: bool = True, num_classes: int = 80, device: str = None) -> nn.Module:
107 | return create_yolox_model("yolov3", pretrained, num_classes, device)
108 |
109 |
110 | def yolox_custom(ckpt_path: str = None, exp_path: str = None, device: str = None) -> nn.Module:
111 | return create_yolox_model("yolox_custom", ckpt_path=ckpt_path, exp_path=exp_path, device=device)
112 |
--------------------------------------------------------------------------------
/tools/train.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Copyright (c) Megvii, Inc. and its affiliates.
4 |
5 | import argparse
6 | import random
7 | import warnings
8 | from loguru import logger
9 |
10 | import torch
11 | import torch.backends.cudnn as cudnn
12 |
13 | from yolox.core import launch
14 | from yolox.exp import Exp, check_exp_value, get_exp
15 | from yolox.utils import configure_module, configure_nccl, configure_omp, get_num_devices
16 |
17 |
18 | def make_parser():
19 | parser = argparse.ArgumentParser("YOLOX train parser")
20 | parser.add_argument("-expn", "--experiment-name", type=str, default=None)
21 | parser.add_argument("-n", "--name", type=str, default=None, help="model name")
22 |
23 | # distributed
24 | parser.add_argument(
25 | "--dist-backend", default="nccl", type=str, help="distributed backend"
26 | )
27 | parser.add_argument(
28 | "--dist-url",
29 | default=None,
30 | type=str,
31 | help="url used to set up distributed training",
32 | )
33 | parser.add_argument("-b", "--batch-size", type=int, default=64, help="batch size")
34 | parser.add_argument(
35 | "-d", "--devices", default=None, type=int, help="device for training"
36 | )
37 | parser.add_argument(
38 | "-f",
39 | "--exp_file",
40 | default=None,
41 | type=str,
42 | help="plz input your experiment description file",
43 | )
44 | parser.add_argument(
45 | "--resume", default=False, action="store_true", help="resume training"
46 | )
47 | parser.add_argument("-c", "--ckpt", default=None, type=str, help="checkpoint file")
48 | parser.add_argument(
49 | "-e",
50 | "--start_epoch",
51 | default=None,
52 | type=int,
53 | help="resume training start epoch",
54 | )
55 | parser.add_argument(
56 | "--num_machines", default=1, type=int, help="num of node for training"
57 | )
58 | parser.add_argument(
59 | "--machine_rank", default=0, type=int, help="node rank for multi-node training"
60 | )
61 | parser.add_argument(
62 | "--fp16",
63 | dest="fp16",
64 | default=False,
65 | action="store_true",
66 | help="Adopting mix precision training.",
67 | )
68 | parser.add_argument(
69 | "--cache",
70 | type=str,
71 | nargs="?",
72 | const="ram",
73 | help="Caching imgs to ram/disk for fast training.",
74 | )
75 | parser.add_argument(
76 | "-o",
77 | "--occupy",
78 | dest="occupy",
79 | default=False,
80 | action="store_true",
81 | help="occupy GPU memory first for training.",
82 | )
83 | parser.add_argument(
84 | "-l",
85 | "--logger",
86 | type=str,
87 | help="Logger to be used for metrics. \
88 | Implemented loggers include `tensorboard`, `mlflow` and `wandb`.",
89 | default="tensorboard"
90 | )
91 | parser.add_argument(
92 | "opts",
93 | help="Modify config options using the command-line",
94 | default=None,
95 | nargs=argparse.REMAINDER,
96 | )
97 | return parser
98 |
99 |
100 | @logger.catch
101 | def main(exp: Exp, args):
102 | if exp.seed is not None:
103 | random.seed(exp.seed)
104 | torch.manual_seed(exp.seed)
105 | cudnn.deterministic = True
106 | warnings.warn(
107 | "You have chosen to seed training. This will turn on the CUDNN deterministic setting, "
108 | "which can slow down your training considerably! You may see unexpected behavior "
109 | "when restarting from checkpoints."
110 | )
111 |
112 | # set environment variables for distributed training
113 | configure_nccl()
114 | configure_omp()
115 | cudnn.benchmark = True
116 |
117 | trainer = exp.get_trainer(args)
118 | trainer.train()
119 |
120 |
121 | if __name__ == "__main__":
122 | configure_module()
123 | args = make_parser().parse_args()
124 | exp = get_exp(args.exp_file, args.name)
125 | exp.merge(args.opts)
126 | check_exp_value(exp)
127 |
128 | if not args.experiment_name:
129 | args.experiment_name = exp.exp_name
130 |
131 | num_gpu = get_num_devices() if args.devices is None else args.devices
132 | assert num_gpu <= get_num_devices()
133 |
134 | if args.cache is not None:
135 | exp.dataset = exp.get_dataset(cache=True, cache_type=args.cache)
136 |
137 | dist_url = "auto" if args.dist_url is None else args.dist_url
138 | launch(
139 | main,
140 | num_gpu,
141 | args.num_machines,
142 | args.machine_rank,
143 | backend=args.dist_backend,
144 | dist_url=dist_url,
145 | args=(exp, args),
146 | )
147 |
--------------------------------------------------------------------------------
/demo/ncnn/cpp/README.md:
--------------------------------------------------------------------------------
1 | # YOLOX-CPP-ncnn
2 |
3 | Cpp file compile of YOLOX object detection base on [ncnn](https://github.com/Tencent/ncnn).
4 |
5 | ## Tutorial
6 |
7 | ### Step1
8 | Clone [ncnn](https://github.com/Tencent/ncnn) first, then please following [build tutorial of ncnn](https://github.com/Tencent/ncnn/wiki/how-to-build) to build on your own device.
9 |
10 | ### Step2
11 | First, we try the original onnx2ncnn solution by using provided tools to generate onnx file.
12 | For example, if you want to generate onnx file of yolox-s, please run the following command:
13 | ```shell
14 | cd
15 | python3 tools/export_onnx.py -n yolox-s
16 | ```
17 | Then a yolox.onnx file is generated.
18 |
19 | ### Step3
20 | Generate ncnn param and bin file.
21 | ```shell
22 | cd
23 | cd build/tools/ncnn
24 | ./onnx2ncnn yolox.onnx model.param model.bin
25 | ```
26 |
27 | Since Focus module is not supported in ncnn. You will see warnings like:
28 | ```shell
29 | Unsupported slice step!
30 | ```
31 | However, don't worry on this as a C++ version of Focus layer is already implemented in yolox.cpp.
32 |
33 | ### Step4
34 | Open **model.param**, and modify it. For more information on the ncnn param and model file structure, please take a look at this [wiki](https://github.com/Tencent/ncnn/wiki/param-and-model-file-structure).
35 |
36 | Before (just an example):
37 | ```
38 | 295 328
39 | Input images 0 1 images
40 | Split splitncnn_input0 1 4 images images_splitncnn_0 images_splitncnn_1 images_splitncnn_2 images_splitncnn_3
41 | Crop Slice_4 1 1 images_splitncnn_3 647 -23309=1,0 -23310=1,2147483647 -23311=1,1
42 | Crop Slice_9 1 1 647 652 -23309=1,0 -23310=1,2147483647 -23311=1,2
43 | Crop Slice_14 1 1 images_splitncnn_2 657 -23309=1,0 -23310=1,2147483647 -23311=1,1
44 | Crop Slice_19 1 1 657 662 -23309=1,1 -23310=1,2147483647 -23311=1,2
45 | Crop Slice_24 1 1 images_splitncnn_1 667 -23309=1,1 -23310=1,2147483647 -23311=1,1
46 | Crop Slice_29 1 1 667 672 -23309=1,0 -23310=1,2147483647 -23311=1,2
47 | Crop Slice_34 1 1 images_splitncnn_0 677 -23309=1,1 -23310=1,2147483647 -23311=1,1
48 | Crop Slice_39 1 1 677 682 -23309=1,1 -23310=1,2147483647 -23311=1,2
49 | Concat Concat_40 4 1 652 672 662 682 683 0=0
50 | ...
51 | ```
52 | * Change first number for 295 to 295 - 9 = 286 (since we will remove 10 layers and add 1 layers, total layers number should minus 9).
53 | * Then remove 10 lines of code from Split to Concat, but remember the last but 2nd number: 683.
54 | * Add YoloV5Focus layer After Input (using previous number 683):
55 | ```
56 | YoloV5Focus focus 1 1 images 683
57 | ```
58 | After(just an example):
59 | ```
60 | 286 328
61 | Input images 0 1 images
62 | YoloV5Focus focus 1 1 images 683
63 | ...
64 | ```
65 |
66 | ### Step5
67 | Use ncnn_optimize to generate new param and bin:
68 | ```shell
69 | # suppose you are still under ncnn/build/tools/ncnn dir.
70 | ../ncnnoptimize model.param model.bin yolox.param yolox.bin 65536
71 | ```
72 |
73 | ### Step6
74 | Copy or Move yolox.cpp file into ncnn/examples, modify the CMakeList.txt to add our implementation, then build.
75 |
76 | ### Step7
77 | Inference image with executable file yolox, enjoy the detect result:
78 | ```shell
79 | ./yolox demo.jpg
80 | ```
81 |
82 | ### Bounus Solution:
83 | As ncnn has released another model conversion tool called [pnnx](https://zhuanlan.zhihu.com/p/427620428) which directly finishs the pytorch2ncnn process via torchscript, we can also try on this.
84 |
85 | ```shell
86 | # take yolox-s as an example
87 | python3 tools/export_torchscript.py -n yolox-s -c /path/to/your_checkpoint_files
88 | ```
89 | Then a `yolox.torchscript.pt` will be generated. Copy this file to your pnnx build directory (pnnx also provides pre-built packages [here](https://github.com/pnnx/pnnx/releases/tag/20220720)).
90 |
91 | ```shell
92 | # suppose you put the yolox.torchscript.pt in a seperate folder
93 | ./pnnx yolox/yolox.torchscript.pt inputshape=[1,3,640,640]
94 | # for zsh users, please use inputshape='[1,3,640,640]'
95 | ```
96 | Still, as ncnn does not support `slice` op as we mentioned in [Step3](https://github.com/Megvii-BaseDetection/YOLOX/tree/main/demo/ncnn/cpp#step3). You will still see the warnings during this process.
97 |
98 | Then multiple pnnx related files will be genreated in your yolox folder. Use `yolox.torchscript.ncnn.param` and `yolox.torchscript.ncnn.bin` as your converted model.
99 |
100 | Then we can follow back to our [Step4](https://github.com/Megvii-BaseDetection/YOLOX/tree/main/demo/ncnn/cpp#step4) for the rest of our implementation.
101 |
102 | ## Acknowledgement
103 |
104 | * [ncnn](https://github.com/Tencent/ncnn)
105 |
--------------------------------------------------------------------------------
/docs/model_zoo.md:
--------------------------------------------------------------------------------
1 | # Model Zoo
2 |
3 | ## Standard Models.
4 |
5 | |Model |size |mAPval
0.5:0.95 |mAPtest
0.5:0.95 | Speed V100
(ms) | Params
(M) |FLOPs
(G)| weights |
6 | | ------ |:---: | :---: | :---: |:---: |:---: | :---: | :----: |
7 | |[YOLOX-s](https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolox_s.py) |640 |40.5 |40.5 |9.8 |9.0 | 26.8 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_s.pth) |
8 | |[YOLOX-m](https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolox_m.py) |640 |46.9 |47.2 |12.3 |25.3 |73.8| [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_m.pth) |
9 | |[YOLOX-l](https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolox_l.py) |640 |49.7 |50.1 |14.5 |54.2| 155.6 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_l.pth) |
10 | |[YOLOX-x](https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolox_x.py) |640 |51.1 |**51.5** | 17.3 |99.1 |281.9 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_x.pth) |
11 | |[YOLOX-Darknet53](https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolov3.py) |640 | 47.7 | 48.0 | 11.1 |63.7 | 185.3 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_darknet.pth)
12 |
13 |
14 | Legacy models
15 |
16 | |Model |size |mAPtest
0.5:0.95 | Speed V100
(ms) | Params
(M) |FLOPs
(G)| weights |
17 | | ------ |:---: | :---: |:---: |:---: | :---: | :----: |
18 | |[YOLOX-s](https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolox_s.py) |640 |39.6 |9.8 |9.0 | 26.8 | [onedrive](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/EW62gmO2vnNNs5npxjzunVwB9p307qqygaCkXdTO88BLUg?e=NMTQYw)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_s.pth) |
19 | |[YOLOX-m](https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolox_m.py) |640 |46.4 |12.3 |25.3 |73.8| [onedrive](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/ERMTP7VFqrVBrXKMU7Vl4TcBQs0SUeCT7kvc-JdIbej4tQ?e=1MDo9y)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_m.pth) |
20 | |[YOLOX-l](https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolox_l.py) |640 |50.0 |14.5 |54.2| 155.6 | [onedrive](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/EWA8w_IEOzBKvuueBqfaZh0BeoG5sVzR-XYbOJO4YlOkRw?e=wHWOBE)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_l.pth) |
21 | |[YOLOX-x](https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolox_x.py) |640 |**51.2** | 17.3 |99.1 |281.9 | [onedrive](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/EdgVPHBziOVBtGAXHfeHI5kBza0q9yyueMGdT0wXZfI1rQ?e=tABO5u)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_x.pth) |
22 | |[YOLOX-Darknet53](https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolov3.py) |640 | 47.4 | 11.1 |63.7 | 185.3 | [onedrive](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/EZ-MV1r_fMFPkPrNjvbJEMoBLOLAnXH-XKEB77w8LhXL6Q?e=mf6wOc)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_darknet53.pth) |
23 |
24 |
25 |
26 | ## Light Models.
27 |
28 | |Model |size |mAPval
0.5:0.95 | Params
(M) |FLOPs
(G)| weights |
29 | | ------ |:---: | :---: |:---: |:---: | :---: |
30 | |[YOLOX-Nano](https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolox_nano.py) |416 |25.8 | 0.91 |1.08 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_nano.pth) |
31 | |[YOLOX-Tiny](https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolox_tiny.py) |416 |32.8 | 5.06 |6.45 | [github](https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_tiny.pth) |
32 |
33 |
34 |
35 | Legacy models
36 |
37 | |Model |size |mAPval
0.5:0.95 | Params
(M) |FLOPs
(G)| weights |
38 | | ------ |:---: | :---: |:---: |:---: | :---: |
39 | |[YOLOX-Nano](https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolox_nano.py) |416 |25.3 | 0.91 |1.08 | [onedrive](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/EdcREey-krhLtdtSnxolxiUBjWMy6EFdiaO9bdOwZ5ygCQ?e=yQpdds)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_nano.pth) |
40 | |[YOLOX-Tiny](https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolox_tiny.py) |416 |32.8 | 5.06 |6.45 | [onedrive](https://megvii-my.sharepoint.cn/:u:/g/personal/gezheng_megvii_com/EbZuinX5X1dJmNy8nqSRegABWspKw3QpXxuO82YSoFN1oQ?e=Q7V7XE)/[github](https://github.com/Megvii-BaseDetection/storage/releases/download/0.0.1/yolox_tiny_32dot8.pth) |
41 |
42 |
43 |
--------------------------------------------------------------------------------
/yolox/core/launch.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # -*- coding:utf-8 -*-
3 | # Code are based on
4 | # https://github.com/facebookresearch/detectron2/blob/master/detectron2/engine/launch.py
5 | # Copyright (c) Facebook, Inc. and its affiliates.
6 | # Copyright (c) Megvii, Inc. and its affiliates.
7 |
8 | import sys
9 | from datetime import timedelta
10 | from loguru import logger
11 |
12 | import torch
13 | import torch.distributed as dist
14 | import torch.multiprocessing as mp
15 |
16 | import yolox.utils.dist as comm
17 |
18 | __all__ = ["launch"]
19 |
20 |
21 | DEFAULT_TIMEOUT = timedelta(minutes=30)
22 |
23 |
24 | def _find_free_port():
25 | """
26 | Find an available port of current machine / node.
27 | """
28 | import socket
29 |
30 | sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
31 | # Binding to port 0 will cause the OS to find an available port for us
32 | sock.bind(("", 0))
33 | port = sock.getsockname()[1]
34 | sock.close()
35 | # NOTE: there is still a chance the port could be taken by other processes.
36 | return port
37 |
38 |
39 | def launch(
40 | main_func,
41 | num_gpus_per_machine,
42 | num_machines=1,
43 | machine_rank=0,
44 | backend="nccl",
45 | dist_url=None,
46 | args=(),
47 | timeout=DEFAULT_TIMEOUT,
48 | ):
49 | """
50 | Args:
51 | main_func: a function that will be called by `main_func(*args)`
52 | num_machines (int): the total number of machines
53 | machine_rank (int): the rank of this machine (one per machine)
54 | dist_url (str): url to connect to for distributed training, including protocol
55 | e.g. "tcp://127.0.0.1:8686".
56 | Can be set to auto to automatically select a free port on localhost
57 | args (tuple): arguments passed to main_func
58 | """
59 | world_size = num_machines * num_gpus_per_machine
60 | if world_size > 1:
61 | # https://github.com/pytorch/pytorch/pull/14391
62 | # TODO prctl in spawned processes
63 |
64 | if dist_url == "auto":
65 | assert (
66 | num_machines == 1
67 | ), "dist_url=auto cannot work with distributed training."
68 | port = _find_free_port()
69 | dist_url = f"tcp://127.0.0.1:{port}"
70 |
71 | start_method = "spawn"
72 | cache = vars(args[1]).get("cache", False)
73 |
74 | # To use numpy memmap for caching image into RAM, we have to use fork method
75 | if cache:
76 | assert sys.platform != "win32", (
77 | "As Windows platform doesn't support fork method, "
78 | "do not add --cache in your training command."
79 | )
80 | start_method = "fork"
81 |
82 | mp.start_processes(
83 | _distributed_worker,
84 | nprocs=num_gpus_per_machine,
85 | args=(
86 | main_func,
87 | world_size,
88 | num_gpus_per_machine,
89 | machine_rank,
90 | backend,
91 | dist_url,
92 | args,
93 | ),
94 | daemon=False,
95 | start_method=start_method,
96 | )
97 | else:
98 | main_func(*args)
99 |
100 |
101 | def _distributed_worker(
102 | local_rank,
103 | main_func,
104 | world_size,
105 | num_gpus_per_machine,
106 | machine_rank,
107 | backend,
108 | dist_url,
109 | args,
110 | timeout=DEFAULT_TIMEOUT,
111 | ):
112 | assert (
113 | torch.cuda.is_available()
114 | ), "cuda is not available. Please check your installation."
115 | global_rank = machine_rank * num_gpus_per_machine + local_rank
116 | logger.info("Rank {} initialization finished.".format(global_rank))
117 | try:
118 | dist.init_process_group(
119 | backend=backend,
120 | init_method=dist_url,
121 | world_size=world_size,
122 | rank=global_rank,
123 | timeout=timeout,
124 | )
125 | except Exception:
126 | logger.error("Process group URL: {}".format(dist_url))
127 | raise
128 |
129 | # Setup the local process group (which contains ranks within the same machine)
130 | assert comm._LOCAL_PROCESS_GROUP is None
131 | num_machines = world_size // num_gpus_per_machine
132 | for i in range(num_machines):
133 | ranks_on_i = list(
134 | range(i * num_gpus_per_machine, (i + 1) * num_gpus_per_machine)
135 | )
136 | pg = dist.new_group(ranks_on_i)
137 | if i == machine_rank:
138 | comm._LOCAL_PROCESS_GROUP = pg
139 |
140 | # synchronize is needed here to prevent a possible timeout after calling init_process_group
141 | # See: https://github.com/facebookresearch/maskrcnn-benchmark/issues/172
142 | comm.synchronize()
143 |
144 | assert num_gpus_per_machine <= torch.cuda.device_count()
145 | torch.cuda.set_device(local_rank)
146 |
147 | main_func(*args)
148 |
--------------------------------------------------------------------------------
/yolox/layers/jit_ops.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # Copyright (c) Megvii, Inc. and its affiliates. All Rights Reserved
3 |
4 | import glob
5 | import importlib
6 | import os
7 | import sys
8 | import time
9 | from typing import List
10 |
11 | __all__ = ["JitOp", "FastCOCOEvalOp"]
12 |
13 |
14 | class JitOp:
15 | """
16 | Just-in-time compilation of ops.
17 |
18 | Some code of `JitOp` is inspired by `deepspeed.op_builder`,
19 | check the following link for more details:
20 | https://github.com/microsoft/DeepSpeed/blob/master/op_builder/builder.py
21 | """
22 |
23 | def __init__(self, name):
24 | self.name = name
25 |
26 | def absolute_name(self) -> str:
27 | """Get absolute build path for cases where the op is pre-installed."""
28 | pass
29 |
30 | def sources(self) -> List:
31 | """Get path list of source files of op.
32 |
33 | NOTE: the path should be elative to root of package during building,
34 | Otherwise, exception will be raised when building package.
35 | However, for runtime building, path will be absolute.
36 | """
37 | pass
38 |
39 | def include_dirs(self) -> List:
40 | """
41 | Get list of include paths, relative to root of package.
42 |
43 | NOTE: the path should be elative to root of package.
44 | Otherwise, exception will be raised when building package.
45 | """
46 | return []
47 |
48 | def define_macros(self) -> List:
49 | """Get list of macros to define for op"""
50 | return []
51 |
52 | def cxx_args(self) -> List:
53 | """Get optional list of compiler flags to forward"""
54 | args = ["-O2"] if sys.platform == "win32" else ["-O3", "-std=c++14", "-g", "-Wno-reorder"]
55 | return args
56 |
57 | def nvcc_args(self) -> List:
58 | """Get optional list of compiler flags to forward to nvcc when building CUDA sources"""
59 | args = [
60 | "-O3", "--use_fast_math",
61 | "-std=c++17" if sys.platform == "win32" else "-std=c++14",
62 | "-U__CUDA_NO_HALF_OPERATORS__",
63 | "-U__CUDA_NO_HALF_CONVERSIONS__",
64 | "-U__CUDA_NO_HALF2_OPERATORS__",
65 | ]
66 | return args
67 |
68 | def build_op(self):
69 | from torch.utils.cpp_extension import CppExtension
70 | return CppExtension(
71 | name=self.absolute_name(),
72 | sources=self.sources(),
73 | include_dirs=self.include_dirs(),
74 | define_macros=self.define_macros(),
75 | extra_compile_args={
76 | "cxx": self.cxx_args(),
77 | },
78 | )
79 |
80 | def load(self, verbose=True):
81 | try:
82 | # try to import op from pre-installed package
83 | return importlib.import_module(self.absolute_name())
84 | except Exception: # op not compiled, jit load
85 | from yolox.utils import wait_for_the_master
86 | with wait_for_the_master(): # to avoid race condition
87 | return self.jit_load(verbose)
88 |
89 | def jit_load(self, verbose=True):
90 | from torch.utils.cpp_extension import load
91 | from loguru import logger
92 | try:
93 | import ninja # noqa
94 | except ImportError:
95 | if verbose:
96 | logger.warning(
97 | f"Ninja is not installed, fall back to normal installation for {self.name}."
98 | )
99 |
100 | build_tik = time.time()
101 | # build op and load
102 | op_module = load(
103 | name=self.name,
104 | sources=self.sources(),
105 | extra_cflags=self.cxx_args(),
106 | extra_cuda_cflags=self.nvcc_args(),
107 | verbose=verbose,
108 | )
109 | build_duration = time.time() - build_tik
110 | if verbose:
111 | logger.info(f"Load {self.name} op in {build_duration:.3f}s.")
112 | return op_module
113 |
114 | def clear_dynamic_library(self):
115 | """Remove dynamic libraray files generated by JIT compilation."""
116 | module = self.load()
117 | os.remove(module.__file__)
118 |
119 |
120 | class FastCOCOEvalOp(JitOp):
121 |
122 | def __init__(self, name="fast_cocoeval"):
123 | super().__init__(name=name)
124 |
125 | def absolute_name(self):
126 | return f'yolox.layers.{self.name}'
127 |
128 | def sources(self):
129 | sources = glob.glob(os.path.join("yolox", "layers", "cocoeval", "*.cpp"))
130 | if not sources: # source will be empty list if the so file is removed after install
131 | # use abosolute path to compile
132 | import yolox
133 | code_path = os.path.join(yolox.__path__[0], "layers", "cocoeval", "*.cpp")
134 | sources = glob.glob(code_path)
135 | return sources
136 |
137 | def include_dirs(self):
138 | return [os.path.join("yolox", "layers", "cocoeval")]
139 |
--------------------------------------------------------------------------------
/docs/cache.md:
--------------------------------------------------------------------------------
1 | # Cache Custom Data
2 |
3 | The caching feature is specifically tailored for users with ample memory resources. However, we still offer the option to cache data to disk, but disk performance can vary and may not guarantee optimal user experience. Implementing custom dataset RAM caching is also more straightforward and user-friendly compared to disk caching. With a few simple modifications, users can expect to see a significant increase in training speed, with speeds nearly double that of non-cached datasets.
4 |
5 | This page explains how to cache your own custom data with YOLOX.
6 |
7 | ## 0. Before you start
8 |
9 | **Step1** Clone this repo and follow the [README](../README.md) to install YOLOX.
10 |
11 | **Stpe2** Read the [Training on custom data](./train_custom_data.md) tutorial to understand how to prepare your custom data.
12 |
13 | ## 1. Inheirit from `CacheDataset`
14 |
15 |
16 | **Step1** Create a custom dataset that inherits from the `CacheDataset` class. Note that whether inheriting from `Dataset` or `CacheDataset `, the `__init__()` method of your custom dataset should take the following keyword arguments: `input_dimension`, `cache`, and `cache_type`. Also, call `super().__init__()` and pass in `input_dimension`, `num_imgs`, `cache`, and `cache_type` as input, where `num_imgs` is the size of the dataset.
17 |
18 | **Step2** Implement the abstract function `read_img(self, index, use_cache=True)` of parent class and decorate it with `@cache_read_img`. This function takes an `index` as input and returns an `image`, and the returned image will be used for caching. It is recommended to put all repetitive and fixed post-processing operations on the image in this function to reduce the post-processing time of the image during training.
19 |
20 | ```python
21 | # CustomDataset.py
22 | from yolox.data.datasets import CacheDataset, cache_read_img
23 |
24 | class CustomDataset(CacheDataset):
25 | def __init__(self, input_dimension, cache, cache_type, *args, **kwargs):
26 | # Get the required keyword arguments of super().__init__()
27 | super().__init__(
28 | input_dimension=input_dimension,
29 | num_imgs=num_imgs,
30 | cache=cache,
31 | cache_type=cache_type
32 | )
33 | # ...
34 |
35 | @cache_read_img
36 | def read_img(self, index, use_cache=True):
37 | # get image ...
38 | # (optional) repetitive and fixed post-processing operations for image
39 | return image
40 | ```
41 |
42 | ## 2. Create your Exp file and return your custom dataset
43 |
44 | **Step1** Create a new class that inherits from the `Exp` class provided by the `yolox_base.py`. Override the `get_dataset()` and `get_eval_dataset()` method to return an instance of your custom dataset.
45 |
46 | **Step2** Implement your own `get_evaluator` method to return an instance of your custom evaluator.
47 |
48 | ```python
49 | # CustomeExp.py
50 | from yolox.exp import Exp as MyExp
51 |
52 | class Exp(MyExp):
53 | def get_dataset(self, cache, cache_type: str = "ram"):
54 | return CustomDataset(
55 | input_dimension=self.input_size,
56 | cache=cache,
57 | cache_type=cache_type
58 | )
59 |
60 | def get_eval_dataset(self):
61 | return CustomDataset(
62 | input_dimension=self.input_size,
63 | )
64 |
65 | def get_evaluator(self, batch_size, is_distributed, testdev=False, legacy=False):
66 | return CustomEvaluator(
67 | dataloader=self.get_eval_loader(batch_size, is_distributed, testdev=testdev, legacy=legacy),
68 | img_size=self.test_size,
69 | confthre=self.test_conf,
70 | nmsthre=self.nmsthre,
71 | num_classes=self.num_classes,
72 | testdev=testdev,
73 | )
74 | ```
75 |
76 | **(Optional)** `get_data_loader` and `get_eval_loader` are now a default behavior in `yolox_base.py` and generally do not need to be changed. If you have to change `get_data_loader`, you need to add the following code at the beginning.
77 |
78 | ```python
79 | # CustomeExp.py
80 | from yolox.exp import Exp as MyExp
81 |
82 | class Exp(MyExp):
83 | def get_data_loader(self, batch_size, is_distributed, no_aug=False, cache_img: str = None):
84 | if self.dataset is None:
85 | with wait_for_the_master():
86 | assert cache_img is None
87 | self.dataset = self.get_dataset(cache=False, cache_type=cache_img)
88 | # ...
89 |
90 | ```
91 |
92 | ## 3. Cache to Disk
93 | It's important to note that the `cache_type` can be `"ram"` or `"disk"`, depending on where you want to cache your dataset. If you choose `"disk"`, you need to pass in additional parameters to `super().__init__()` of `CustomDataset`: `data_dir`, `cache_dir_name`, `path_filename`.
94 |
95 | - `data_dir`: the root directory of the dataset, e.g. `/path/to/COCO`.
96 | - `cache_dir_name`: the name of the directory to cache to disk, for example `"custom_cache"`, then the files cached to disk will be saved under `/path/to/COCO/custom_cache`.
97 | - `path_filename`: a list of paths to the data relative to the `data_dir`, e.g. if you have data `/path/to/COCO/train/1.jpg`, `/path/to/COCO/train/2.jpg`, then `path_filename = ['train/1.jpg', ' train/2.jpg']`.
98 |
--------------------------------------------------------------------------------
/demo/nebullvm/README.md:
--------------------------------------------------------------------------------
1 | # **Accelerate YOLOX inference with nebullvm in Python**
2 |
3 | This document shows how to accelerate YOLOX inference time with nebullvm.
4 |
5 | [nebullvm](https://github.com/nebuly-ai/nebullvm) is an open-source library designed to accelerate AI inference of deep learning models in a few lines of code. nebullvm leverages state-of-the-art model optimization techniques such as deep learning compilers (TensorRT, Openvino, ONNX Runtime, TVM, TF Lite, DeepSparse, etc.), various quantization and compression strategies to achieve the maximum physically possible acceleration on the user's hardware.
6 |
7 | ## Benchmarks
8 | Following are the results of the nebullvm optimization on YOLOX without loss of accuracy.
9 | For each model-hardware pairing, response time was evaluated as the average over 100 predictions. The test was run on Nvidia Tesla T4 (g4dn.xlarge) and Intel XEON Scalable (m6i.24xlarge and c6i.12xlarge) on AWS.
10 |
11 | | Model | Hardware | Unoptimized (ms)| Nebullvm optimized (ms) | Speedup |
12 | |---------|--------------|-----------------|-------------------------|---------|
13 | | YOLOX-s | g4dn.xlarge | 13.6 | 9.0 | 1.5x |
14 | | YOLOX-s | m6i.24xlarge | 32.7 | 8.8 | 3.7x |
15 | | YOLOX-s | c6i.12xlarge | 34.4 | 12.4 | 2.8x |
16 | | YOLOX-m | g4dn.xlarge | 24.2 | 22.4 | 1.1x |
17 | | YOLOX-m | m6i.24xlarge | 55.1 | 36.0 | 2.3x |
18 | | YOLOX-m | c6i.12xlarge | 62.5 | 26.9 | 2.6x |
19 | | YOLOX-l | g4dn.xlarge | 84.4 | 80.5 | 1.5x |
20 | | YOLOX-l | m6i.24xlarge | 88.0 | 33.7 | 2.6x |
21 | | YOLOX-l | c6i.12xlarge | 102.8 | 54.2 | 1.9x |
22 | | YOLOX-x | g4dn.xlarge | 87.3 | 34.0 | 2.6x |
23 | | YOLOX-x | m6i.24xlarge | 134.5 | 56.6 | 2.4x |
24 | | YOLOX-x | c6i.12xlarge | 162.0 | 95.4 | 1.7x |
25 |
26 | ## Steps to accelerate YOLOX with nebullvm
27 | 1. Download a YOLOX model from the original [readme](https://github.com/Megvii-BaseDetection/YOLOX)
28 | 2. Optimize YOLOX with nebullvm
29 | 3. Perform inference and compare the latency of the optimized model with that of the original model
30 |
31 | [Here](nebullvm_optimization.py) you can find a demo in python.
32 |
33 |
34 | First, let's install nebullvm. The simplest way is by using pip.
35 | ```
36 | pip install nebullvm
37 | ```
38 | Now, let's download one of YOLOX models and optimize it with nebullvm.
39 |
40 | ```python
41 | # Import YOLOX model
42 | from yolox.exp import get_exp
43 | from yolox.data.data_augment import ValTransform
44 |
45 | exp = get_exp(None, 'yolox-s') # select model name
46 | model = exp.get_model()
47 | model.cuda()
48 | model.eval()
49 |
50 |
51 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
52 | input_data = [((torch.randn(1, 3, 640, 640).to(device), ), 0) for i in range(100)]
53 |
54 | # Run nebullvm optimization without performance loss
55 | optimized_model = optimize_model(model, input_data=input_data, optimization_time="constrained")
56 | ```
57 | Find [here](nebullvm_optimize.py) the complete script in python with more details.
58 |
59 | In this example, we optimized YOLOX without any loss in accuracy. To further speed up the model by means of more aggressive optimization techniques, proceed as follows:
60 | - Set *optimization_time="unconstrained"*. With the unconstrained option, nebullvm will test time-consuming techniques such as pruning and quantization-aware training (QAT).
61 | - Set the *metric_drop_ths* parameter to be greater than zero (by default, *metric_drop_ths=0*). In this way, we will allow nebullvm to test optimization techniques that involve a tradeoff of some trade-off of a certain metric. For example, to test maximum acceleration with a minimum loss of accuracy of 3%, set *metric_drop_ths=0.03* and *metric="accuracy"*.
62 | For more information about nebullvm API, see [nebullvm documentation](https://github.com/nebuly-ai/nebullvm).
63 |
64 |
65 | Let's now compare the latency of the optimized model with that of the original model.
66 | Note that before testing latency of the optimized model, it is necessary to perform some warmup runs, as some optimizers fine-tune certain internal parameters during the first few inferences after optimization.
67 |
68 | ```python
69 | # Check perfomance
70 | warmup_iters = 30
71 | num_iters = 100
72 |
73 | # Unoptimized model perfomance
74 | with torch.no_grad():
75 | for i in range(warmup_iters):
76 | o = model(img)
77 |
78 | start = time.time()
79 | for i in range(num_iters):
80 | o = model(img)
81 | stop = time.time()
82 | print(f"Average inference time of unoptimized YOLOX: {(stop - start)/num_iters*1000} ms")
83 |
84 | # Optimized model perfomance
85 | with torch.no_grad():
86 | for i in range(warmup_iters):
87 | res = model_opt(img)
88 |
89 | start = time.time()
90 | for i in range(num_iters):
91 | res = model_opt(img)
92 | stop = time.time()
93 | print(f"Average inference time of YOLOX otpimized with nebullvm: {(stop - start)/num_iters*1000} ms")
94 | ```
95 | Find [here](nebullvm_optimization.py) the complete script in python with more details.
96 |
--------------------------------------------------------------------------------
/yolox/utils/boxes.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # Copyright (c) Megvii Inc. All rights reserved.
3 |
4 | import numpy as np
5 |
6 | import torch
7 | import torchvision
8 |
9 | __all__ = [
10 | "filter_box",
11 | "postprocess",
12 | "bboxes_iou",
13 | "matrix_iou",
14 | "adjust_box_anns",
15 | "xyxy2xywh",
16 | "xyxy2cxcywh",
17 | "cxcywh2xyxy",
18 | ]
19 |
20 |
21 | def filter_box(output, scale_range):
22 | """
23 | output: (N, 5+class) shape
24 | """
25 | min_scale, max_scale = scale_range
26 | w = output[:, 2] - output[:, 0]
27 | h = output[:, 3] - output[:, 1]
28 | keep = (w * h > min_scale * min_scale) & (w * h < max_scale * max_scale)
29 | return output[keep]
30 |
31 |
32 | def postprocess(prediction, num_classes, conf_thre=0.7, nms_thre=0.45, class_agnostic=False):
33 | box_corner = prediction.new(prediction.shape)
34 | box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2
35 | box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2
36 | box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2
37 | box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2
38 | prediction[:, :, :4] = box_corner[:, :, :4]
39 |
40 | output = [None for _ in range(len(prediction))]
41 | for i, image_pred in enumerate(prediction):
42 |
43 | # If none are remaining => process next image
44 | if not image_pred.size(0):
45 | continue
46 | # Get score and class with highest confidence
47 | class_conf, class_pred = torch.max(image_pred[:, 5: 5 + num_classes], 1, keepdim=True)
48 |
49 | conf_mask = (image_pred[:, 4] * class_conf.squeeze() >= conf_thre).squeeze()
50 | # Detections ordered as (x1, y1, x2, y2, obj_conf, class_conf, class_pred)
51 | detections = torch.cat((image_pred[:, :5], class_conf, class_pred.float()), 1)
52 | detections = detections[conf_mask]
53 | if not detections.size(0):
54 | continue
55 |
56 | if class_agnostic:
57 | nms_out_index = torchvision.ops.nms(
58 | detections[:, :4],
59 | detections[:, 4] * detections[:, 5],
60 | nms_thre,
61 | )
62 | else:
63 | nms_out_index = torchvision.ops.batched_nms(
64 | detections[:, :4],
65 | detections[:, 4] * detections[:, 5],
66 | detections[:, 6],
67 | nms_thre,
68 | )
69 |
70 | detections = detections[nms_out_index]
71 | if output[i] is None:
72 | output[i] = detections
73 | else:
74 | output[i] = torch.cat((output[i], detections))
75 |
76 | return output
77 |
78 |
79 | def bboxes_iou(bboxes_a, bboxes_b, xyxy=True):
80 | if bboxes_a.shape[1] != 4 or bboxes_b.shape[1] != 4:
81 | raise IndexError
82 |
83 | if xyxy:
84 | tl = torch.max(bboxes_a[:, None, :2], bboxes_b[:, :2])
85 | br = torch.min(bboxes_a[:, None, 2:], bboxes_b[:, 2:])
86 | area_a = torch.prod(bboxes_a[:, 2:] - bboxes_a[:, :2], 1)
87 | area_b = torch.prod(bboxes_b[:, 2:] - bboxes_b[:, :2], 1)
88 | else:
89 | tl = torch.max(
90 | (bboxes_a[:, None, :2] - bboxes_a[:, None, 2:] / 2),
91 | (bboxes_b[:, :2] - bboxes_b[:, 2:] / 2),
92 | )
93 | br = torch.min(
94 | (bboxes_a[:, None, :2] + bboxes_a[:, None, 2:] / 2),
95 | (bboxes_b[:, :2] + bboxes_b[:, 2:] / 2),
96 | )
97 |
98 | area_a = torch.prod(bboxes_a[:, 2:], 1)
99 | area_b = torch.prod(bboxes_b[:, 2:], 1)
100 | en = (tl < br).type(tl.type()).prod(dim=2)
101 | area_i = torch.prod(br - tl, 2) * en # * ((tl < br).all())
102 | return area_i / (area_a[:, None] + area_b - area_i)
103 |
104 |
105 | def matrix_iou(a, b):
106 | """
107 | return iou of a and b, numpy version for data augenmentation
108 | """
109 | lt = np.maximum(a[:, np.newaxis, :2], b[:, :2])
110 | rb = np.minimum(a[:, np.newaxis, 2:], b[:, 2:])
111 |
112 | area_i = np.prod(rb - lt, axis=2) * (lt < rb).all(axis=2)
113 | area_a = np.prod(a[:, 2:] - a[:, :2], axis=1)
114 | area_b = np.prod(b[:, 2:] - b[:, :2], axis=1)
115 | return area_i / (area_a[:, np.newaxis] + area_b - area_i + 1e-12)
116 |
117 |
118 | def adjust_box_anns(bbox, scale_ratio, padw, padh, w_max, h_max):
119 | bbox[:, 0::2] = np.clip(bbox[:, 0::2] * scale_ratio + padw, 0, w_max)
120 | bbox[:, 1::2] = np.clip(bbox[:, 1::2] * scale_ratio + padh, 0, h_max)
121 | return bbox
122 |
123 |
124 | def xyxy2xywh(bboxes):
125 | bboxes[:, 2] = bboxes[:, 2] - bboxes[:, 0]
126 | bboxes[:, 3] = bboxes[:, 3] - bboxes[:, 1]
127 | return bboxes
128 |
129 |
130 | def xyxy2cxcywh(bboxes):
131 | bboxes[:, 2] = bboxes[:, 2] - bboxes[:, 0]
132 | bboxes[:, 3] = bboxes[:, 3] - bboxes[:, 1]
133 | bboxes[:, 0] = bboxes[:, 0] + bboxes[:, 2] * 0.5
134 | bboxes[:, 1] = bboxes[:, 1] + bboxes[:, 3] * 0.5
135 | return bboxes
136 |
137 |
138 | def cxcywh2xyxy(bboxes):
139 | bboxes[:, 0] = bboxes[:, 0] - bboxes[:, 2] * 0.5
140 | bboxes[:, 1] = bboxes[:, 1] - bboxes[:, 3] * 0.5
141 | bboxes[:, 2] = bboxes[:, 0] + bboxes[:, 2]
142 | bboxes[:, 3] = bboxes[:, 1] + bboxes[:, 3]
143 | return bboxes
144 |
--------------------------------------------------------------------------------
/yolox/utils/demo_utils.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | # Copyright (c) Megvii Inc. All rights reserved.
3 |
4 | import os
5 | import random
6 |
7 | import cv2
8 | import numpy as np
9 |
10 | __all__ = [
11 | "mkdir", "nms", "multiclass_nms", "demo_postprocess", "random_color", "visualize_assign"
12 | ]
13 |
14 |
15 | def random_color():
16 | return random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)
17 |
18 |
19 | def visualize_assign(img, boxes, coords, match_results, save_name=None) -> np.ndarray:
20 | """visualize label assign result.
21 |
22 | Args:
23 | img: img to visualize
24 | boxes: gt boxes in xyxy format
25 | coords: coords of matched anchors
26 | match_results: match results of each gt box and coord.
27 | save_name: name of save image, if None, image will not be saved. Default: None.
28 | """
29 | for box_id, box in enumerate(boxes):
30 | x1, y1, x2, y2 = box
31 | color = random_color()
32 | assign_coords = coords[match_results == box_id]
33 | if assign_coords.numel() == 0:
34 | # unmatched boxes are red
35 | color = (0, 0, 255)
36 | cv2.putText(
37 | img, "unmatched", (int(x1), int(y1) - 5),
38 | cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 1
39 | )
40 | else:
41 | for coord in assign_coords:
42 | # draw assigned anchor
43 | cv2.circle(img, (int(coord[0]), int(coord[1])), 3, color, -1)
44 | cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), color, 2)
45 |
46 | if save_name is not None:
47 | cv2.imwrite(save_name, img)
48 |
49 | return img
50 |
51 |
52 | def mkdir(path):
53 | if not os.path.exists(path):
54 | os.makedirs(path)
55 |
56 |
57 | def nms(boxes, scores, nms_thr):
58 | """Single class NMS implemented in Numpy."""
59 | x1 = boxes[:, 0]
60 | y1 = boxes[:, 1]
61 | x2 = boxes[:, 2]
62 | y2 = boxes[:, 3]
63 |
64 | areas = (x2 - x1 + 1) * (y2 - y1 + 1)
65 | order = scores.argsort()[::-1]
66 |
67 | keep = []
68 | while order.size > 0:
69 | i = order[0]
70 | keep.append(i)
71 | xx1 = np.maximum(x1[i], x1[order[1:]])
72 | yy1 = np.maximum(y1[i], y1[order[1:]])
73 | xx2 = np.minimum(x2[i], x2[order[1:]])
74 | yy2 = np.minimum(y2[i], y2[order[1:]])
75 |
76 | w = np.maximum(0.0, xx2 - xx1 + 1)
77 | h = np.maximum(0.0, yy2 - yy1 + 1)
78 | inter = w * h
79 | ovr = inter / (areas[i] + areas[order[1:]] - inter)
80 |
81 | inds = np.where(ovr <= nms_thr)[0]
82 | order = order[inds + 1]
83 |
84 | return keep
85 |
86 |
87 | def multiclass_nms(boxes, scores, nms_thr, score_thr, class_agnostic=True):
88 | """Multiclass NMS implemented in Numpy"""
89 | if class_agnostic:
90 | nms_method = multiclass_nms_class_agnostic
91 | else:
92 | nms_method = multiclass_nms_class_aware
93 | return nms_method(boxes, scores, nms_thr, score_thr)
94 |
95 |
96 | def multiclass_nms_class_aware(boxes, scores, nms_thr, score_thr):
97 | """Multiclass NMS implemented in Numpy. Class-aware version."""
98 | final_dets = []
99 | num_classes = scores.shape[1]
100 | for cls_ind in range(num_classes):
101 | cls_scores = scores[:, cls_ind]
102 | valid_score_mask = cls_scores > score_thr
103 | if valid_score_mask.sum() == 0:
104 | continue
105 | else:
106 | valid_scores = cls_scores[valid_score_mask]
107 | valid_boxes = boxes[valid_score_mask]
108 | keep = nms(valid_boxes, valid_scores, nms_thr)
109 | if len(keep) > 0:
110 | cls_inds = np.ones((len(keep), 1)) * cls_ind
111 | dets = np.concatenate(
112 | [valid_boxes[keep], valid_scores[keep, None], cls_inds], 1
113 | )
114 | final_dets.append(dets)
115 | if len(final_dets) == 0:
116 | return None
117 | return np.concatenate(final_dets, 0)
118 |
119 |
120 | def multiclass_nms_class_agnostic(boxes, scores, nms_thr, score_thr):
121 | """Multiclass NMS implemented in Numpy. Class-agnostic version."""
122 | cls_inds = scores.argmax(1)
123 | cls_scores = scores[np.arange(len(cls_inds)), cls_inds]
124 |
125 | valid_score_mask = cls_scores > score_thr
126 | if valid_score_mask.sum() == 0:
127 | return None
128 | valid_scores = cls_scores[valid_score_mask]
129 | valid_boxes = boxes[valid_score_mask]
130 | valid_cls_inds = cls_inds[valid_score_mask]
131 | keep = nms(valid_boxes, valid_scores, nms_thr)
132 | if keep:
133 | dets = np.concatenate(
134 | [valid_boxes[keep], valid_scores[keep, None], valid_cls_inds[keep, None]], 1
135 | )
136 | return dets
137 |
138 |
139 | def demo_postprocess(outputs, img_size, p6=False):
140 | grids = []
141 | expanded_strides = []
142 | strides = [8, 16, 32] if not p6 else [8, 16, 32, 64]
143 |
144 | hsizes = [img_size[0] // stride for stride in strides]
145 | wsizes = [img_size[1] // stride for stride in strides]
146 |
147 | for hsize, wsize, stride in zip(hsizes, wsizes, strides):
148 | xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize))
149 | grid = np.stack((xv, yv), 2).reshape(1, -1, 2)
150 | grids.append(grid)
151 | shape = grid.shape[:2]
152 | expanded_strides.append(np.full((*shape, 1), stride))
153 |
154 | grids = np.concatenate(grids, 1)
155 | expanded_strides = np.concatenate(expanded_strides, 1)
156 | outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides
157 | outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides
158 |
159 | return outputs
160 |
--------------------------------------------------------------------------------