README.md 4.45 KB
Newer Older
zhe chen's avatar
zhe chen committed
1
2
# InternImage for Semantic Segmentation

zhe chen's avatar
zhe chen committed
3
This folder contains the implementation of the InternImage for semantic segmentation.
zhe chen's avatar
zhe chen committed
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

Our segmentation code is developed on top of [MMSegmentation v0.27.0](https://github.com/open-mmlab/mmsegmentation/tree/v0.27.0).

## Usage

### Install

- Clone this repo:

```bash
git clone https://github.com/OpenGVLab/InternImage.git
cd InternImage
```

- Create a conda virtual environment and activate it:

```bash
conda create -n internimage python=3.7 -y
conda activate internimage
```

- Install `CUDA>=10.2` with `cudnn>=7` following
  the [official installation instructions](https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html)
27
- Install `PyTorch>=1.10.0` and `torchvision>=0.9.0` with `CUDA>=10.2`:
zhe chen's avatar
zhe chen committed
28

29
For examples, to install torch==1.11 with CUDA==11.3 and nvcc:
zhe chen's avatar
zhe chen committed
30

zhe chen's avatar
zhe chen committed
31
```bash
32
33
conda install pytorch==1.11.0 torchvision==0.12.0 torchaudio==0.11.0 cudatoolkit=11.3 -c pytorch -y
conda install -c conda-forge cudatoolkit-dev=11.3 -y # to install nvcc
zhe chen's avatar
zhe chen committed
34
35
```

36
37
- Install other requirements:

zhe chen's avatar
zhe chen committed
38
  note: conda opencv will break torchvision as not to support GPU, so we need to install opencv using pip.
zhe chen's avatar
zhe chen committed
39
40

```bash
41
42
conda install -c conda-forge termcolor yacs pyyaml scipy pip -y
pip install opencv-python
zhe chen's avatar
zhe chen committed
43
44
```

zhe chen's avatar
zhe chen committed
45
- Install `timm` and `mmcv-full` and \`mmsegmentation':
zhe chen's avatar
zhe chen committed
46
47

```bash
48
49
50
51
pip install -U openmim
mim install mmcv-full==1.5.0
mim install mmsegmentation==0.27.0
pip install timm==0.6.11 mmdet==2.28.1
zhe chen's avatar
zhe chen committed
52
53
54
```

- Compile CUDA operators
zhe chen's avatar
zhe chen committed
55

zhe chen's avatar
zhe chen committed
56
57
58
59
60
61
```bash
cd ./ops_dcnv3
sh ./make.sh
# unit test (should see all checking is True)
python test.py
```
zhe chen's avatar
zhe chen committed
62

yeshenglong1's avatar
yeshenglong1 committed
63
- You can also install the operator using .whl files
zhe chen's avatar
zhe chen committed
64
  [DCNv3-1.0-whl](https://github.com/OpenGVLab/InternImage/releases/tag/whl_files)
zhe chen's avatar
zhe chen committed
65

66
### Data Preparation
zhe chen's avatar
zhe chen committed
67

68
Prepare datasets according to the [guidelines](https://github.com/open-mmlab/mmsegmentation/blob/master/docs/en/dataset_prepare.md#prepare-datasets) in MMSegmentation.
zhe chen's avatar
zhe chen committed
69
70
71
72
73
74
75
76

### Evaluation

To evaluate our `InternImage` on ADE20K val, run:

```bash
sh dist_test.sh <config-file> <checkpoint> <gpu-num> --eval mIoU
```
zhe chen's avatar
zhe chen committed
77

78
You can download checkpoint files from [here](https://huggingface.co/OpenGVLab/InternImage/tree/fc1e4e7e01c3e7a39a3875bdebb6577a7256ff91). Then place it to segmentation/checkpoint_dir/seg.
zhe chen's avatar
zhe chen committed
79
80
81
82

For example, to evaluate the `InternImage-T` with a single GPU:

```bash
83
python test.py configs/ade20k/upernet_internimage_t_512_160k_ade20k.py checkpoint_dir/seg/upernet_internimage_t_512_160k_ade20k.pth --eval mIoU
zhe chen's avatar
zhe chen committed
84
85
86
87
88
```

For example, to evaluate the `InternImage-B` with a single node with 8 GPUs:

```bash
89
sh dist_test.sh configs/ade20k/upernet_internimage_b_512_160k_ade20k.py checkpoint_dir/seg/upernet_internimage_b_512_160k_ade20k.pth 8 --eval mIoU
zhe chen's avatar
zhe chen committed
90
91
```

92
### Training
zhe chen's avatar
zhe chen committed
93
94
95
96
97
98
99

To train an `InternImage` on ADE20K, run:

```bash
sh dist_train.sh <config-file> <gpu-num>
```

100
For example, to train `InternImage-T` with 8 GPU on 1 node (total batch size 16), run:
zhe chen's avatar
zhe chen committed
101
102

```bash
103
sh dist_train.sh configs/ade20k/upernet_internimage_t_512_160k_ade20k.py 8
zhe chen's avatar
zhe chen committed
104
105
```

106
### Manage Jobs with Slurm
zhe chen's avatar
zhe chen committed
107

108
For example, to train `InternImage-XL` with 8 GPU on 1 node (total batch size 16), run:
zhe chen's avatar
zhe chen committed
109
110

```bash
111
112
113
114
GPUS=8 sh slurm_train.sh <partition> <job-name> configs/ade20k/upernet_internimage_xl_640_160k_ade20k.py
```

### Image Demo
zhe chen's avatar
zhe chen committed
115

116
117
To inference a single/multiple image like this.
If you specify image containing directory instead of a single image, it will process all the images in the directory.:
zhe chen's avatar
zhe chen committed
118

119
120
121
122
123
```
CUDA_VISIBLE_DEVICES=0 python image_demo.py \
  data/ade/ADEChallengeData2016/images/validation/ADE_val_00000591.jpg \
  configs/ade20k/upernet_internimage_t_512_160k_ade20k.py  \
  checkpoint_dir/seg/upernet_internimage_t_512_160k_ade20k.pth  \
zhe chen's avatar
zhe chen committed
124
  --palette ade20k
zhe chen's avatar
zhe chen committed
125
```
Weiyun1025's avatar
Weiyun1025 committed
126
127
128
129

### Export

To export a segmentation model from PyTorch to TensorRT, run:
zhe chen's avatar
zhe chen committed
130

Weiyun1025's avatar
Weiyun1025 committed
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
```shell
MODEL="model_name"
CKPT_PATH="/path/to/model/ckpt.pth"

python deploy.py \
    "./deploy/configs/mmseg/segmentation_tensorrt_static-512x512.py" \
    "./configs/ade20k/${MODEL}.py" \
    "${CKPT_PATH}" \
    "./deploy/demo.png" \
    --work-dir "./work_dirs/mmseg/${MODEL}" \
    --device cuda \
    --dump-info
```

For example, to export `upernet_internimage_t_512_160k_ade20k` from PyTorch to TensorRT, run:
zhe chen's avatar
zhe chen committed
146

Weiyun1025's avatar
Weiyun1025 committed
147
148
149
150
151
152
153
154
155
156
157
158
159
```shell
MODEL="upernet_internimage_t_512_160k_ade20k"
CKPT_PATH="/path/to/model/ckpt/upernet_internimage_t_512_160k_ade20k.pth"

python deploy.py \
    "./deploy/configs/mmseg/segmentation_tensorrt_static-512x512.py" \
    "./configs/ade20k/${MODEL}.py" \
    "${CKPT_PATH}" \
    "./deploy/demo.png" \
    --work-dir "./work_dirs/mmseg/${MODEL}" \
    --device cuda \
    --dump-info
```