[Tools] Add vit/swin/mit convert weight scripts (#783)

* init scripts * update markdown * update markdown * add docs * delete mit converter and use torch load function * rename segformer readme * update doc * modify doc * 更新中文文档 * Update useful_tools.md * Update useful_tools.md * modify doc * update segformer.yml
open-mmlab · Aug 18, 2021 · f72727c · f72727c
1 parent 441be4e
commit f72727c
Show file tree

Hide file tree

Showing 8 changed files with 381 additions and 11 deletions.
diff --git a/configs/segformer/readme.md → configs/segformer/README.md b/configs/segformer/readme.md → configs/segformer/README.md
@@ -29,15 +29,15 @@
 
 Evaluation with AlignedResize:
 
-| Method | Backbone | Crop Size | Lr schd | mIoU | mIoU(ms+flip) |
-| ------ | -------- | --------- | ------: | ---: | ------------- |
-|Segformer | MIT-B0 | 512x512 | 160000 | 38.1 | 38.57 |
-|Segformer | MIT-B1 | 512x512 | 160000 | 41.64 | 42.76 |
-|Segformer | MIT-B2 | 512x512 | 160000 | 46.53 | 47.49 |
-|Segformer | MIT-B3 | 512x512 | 160000 | 48.46 | 49.14 |
-|Segformer | MIT-B4 | 512x512 | 160000 | 49.34 | 50.29 |
-|Segformer | MIT-B5 | 512x512 | 160000 | 50.08 | 50.72 |
-|Segformer | MIT-B5 | 640x640 | 160000 | 50.58 | 50.8 |
+ | Method | Backbone | Crop Size | Lr schd | mIoU | mIoU(ms+flip) |
+ | ------ | -------- | --------- | ------: | ---: | ------------- |
+ |Segformer | MIT-B0 | 512x512 | 160000 | 38.1 | 38.57 |
+ |Segformer | MIT-B1 | 512x512 | 160000 | 41.64 | 42.76 |
+ |Segformer | MIT-B2 | 512x512 | 160000 | 46.53 | 47.49 |
+ |Segformer | MIT-B3 | 512x512 | 160000 | 48.46 | 49.14 |
+ |Segformer | MIT-B4 | 512x512 | 160000 | 49.34 | 50.29 |
+ |Segformer | MIT-B5 | 512x512 | 160000 | 50.08 | 50.72 |
+ |Segformer | MIT-B5 | 640x640 | 160000 | 50.58 | 50.8 |
 
 We replace `AlignedResize` in original implementatiuon to `Resize + ResizeToMultiple`. If you want to test by
 using `AlignedResize`, you can change the dataset pipeline like this:

diff --git a/configs/segformer/segformer.yml b/configs/segformer/segformer.yml
@@ -0,0 +1,160 @@
+Collections:
+- Metadata:
+ Training Data:
+ - ADE20k
+ Name: segformer
+Models:
+- Config: configs/segformer/segformer_mit-b0_512x512_160k_ade20k.py
+ In Collection: segformer
+ Metadata:
+ backbone: MIT-B0
+ crop size: (512,512)
+ inference time (ms/im):
+ - backend: PyTorch
+ batch size: 1
+ hardware: V100
+ mode: FP32
+ resolution: (512,512)
+ value: 19.49
+ lr schd: 160000
+ memory (GB): 2.1
+ Name: segformer_mit-b0_512x512_160k_ade20k
+ Results:
+ Dataset: ADE20k
+ Metrics:
+ mIoU: 37.41
+ mIoU(ms+flip): 38.34
+ Task: Semantic Segmentation
+ Weights: https://download.openmmlab.com/mmsegmentation/v0.5/segformer/segformer_mit-b0_512x512_160k_ade20k/segformer_mit-b0_512x512_160k_ade20k_20210726_101530-8ffa8fda.pth
+- Config: configs/segformer/segformer_mit-b1_512x512_160k_ade20k.py
+ In Collection: segformer
+ Metadata:
+ backbone: MIT-B1
+ crop size: (512,512)
+ inference time (ms/im):
+ - backend: PyTorch
+ batch size: 1
+ hardware: V100
+ mode: FP32
+ resolution: (512,512)
+ value: 20.98
+ lr schd: 160000
+ memory (GB): 2.6
+ Name: segformer_mit-b1_512x512_160k_ade20k
+ Results:
+ Dataset: ADE20k
+ Metrics:
+ mIoU: 40.97
+ mIoU(ms+flip): 42.54
+ Task: Semantic Segmentation
+ Weights: https://download.openmmlab.com/mmsegmentation/v0.5/segformer/segformer_mit-b1_512x512_160k_ade20k/segformer_mit-b1_512x512_160k_ade20k_20210726_112106-d70e859d.pth
+- Config: configs/segformer/segformer_mit-b2_512x512_160k_ade20k.py
+ In Collection: segformer
+ Metadata:
+ backbone: MIT-B2
+ crop size: (512,512)
+ inference time (ms/im):
+ - backend: PyTorch
+ batch size: 1
+ hardware: V100
+ mode: FP32
+ resolution: (512,512)
+ value: 32.38
+ lr schd: 160000
+ memory (GB): 3.6
+ Name: segformer_mit-b2_512x512_160k_ade20k
+ Results:
+ Dataset: ADE20k
+ Metrics:
+ mIoU: 45.58
+ mIoU(ms+flip): 47.03
+ Task: Semantic Segmentation
+ Weights: https://download.openmmlab.com/mmsegmentation/v0.5/segformer/segformer_mit-b2_512x512_160k_ade20k/segformer_mit-b2_512x512_160k_ade20k_20210726_112103-cbd414ac.pth
+- Config: configs/segformer/segformer_mit-b3_512x512_160k_ade20k.py
+ In Collection: segformer
+ Metadata:
+ backbone: MIT-B3
+ crop size: (512,512)
+ inference time (ms/im):
+ - backend: PyTorch
+ batch size: 1
+ hardware: V100
+ mode: FP32
+ resolution: (512,512)
+ value: 45.23
+ lr schd: 160000
+ memory (GB): 4.8
+ Name: segformer_mit-b3_512x512_160k_ade20k
+ Results:
+ Dataset: ADE20k
+ Metrics:
+ mIoU: 47.82
+ mIoU(ms+flip): 48.81
+ Task: Semantic Segmentation
+ Weights: https://download.openmmlab.com/mmsegmentation/v0.5/segformer/segformer_mit-b3_512x512_160k_ade20k/segformer_mit-b3_512x512_160k_ade20k_20210726_081410-962b98d2.pth
+- Config: configs/segformer/segformer_mit-b4_512x512_160k_ade20k.py
+ In Collection: segformer
+ Metadata:
+ backbone: MIT-B4
+ crop size: (512,512)
+ inference time (ms/im):
+ - backend: PyTorch
+ batch size: 1
+ hardware: V100
+ mode: FP32
+ resolution: (512,512)
+ value: 64.72
+ lr schd: 160000
+ memory (GB): 6.1
+ Name: segformer_mit-b4_512x512_160k_ade20k
+ Results:
+ Dataset: ADE20k
+ Metrics:
+ mIoU: 48.46
+ mIoU(ms+flip): 49.76
+ Task: Semantic Segmentation
+ Weights: https://download.openmmlab.com/mmsegmentation/v0.5/segformer/segformer_mit-b4_512x512_160k_ade20k/segformer_mit-b4_512x512_160k_ade20k_20210728_183055-7f509d7d.pth
+- Config: configs/segformer/segformer_mit-b5_512x512_160k_ade20k.py
+ In Collection: segformer
+ Metadata:
+ backbone: MIT-B5
+ crop size: (512,512)
+ inference time (ms/im):
+ - backend: PyTorch
+ batch size: 1
+ hardware: V100
+ mode: FP32
+ resolution: (512,512)
+ value: 84.1
+ lr schd: 160000
+ memory (GB): 7.2
+ Name: segformer_mit-b5_512x512_160k_ade20k
+ Results:
+ Dataset: ADE20k
+ Metrics:
+ mIoU: 49.13
+ mIoU(ms+flip): 50.22
+ Task: Semantic Segmentation
+ Weights: https://download.openmmlab.com/mmsegmentation/v0.5/segformer/segformer_mit-b5_512x512_160k_ade20k/segformer_mit-b5_512x512_160k_ade20k_20210726_145235-94cedf59.pth
+- Config: configs/segformer/segformer_mit-b5_640x640_160k_ade20k.py
+ In Collection: segformer
+ Metadata:
+ backbone: MIT-B5
+ crop size: (640,640)
+ inference time (ms/im):
+ - backend: PyTorch
+ batch size: 1
+ hardware: V100
+ mode: FP32
+ resolution: (640,640)
+ value: 88.5
+ lr schd: 160000
+ memory (GB): 11.5
+ Name: segformer_mit-b5_640x640_160k_ade20k
+ Results:
+ Dataset: ADE20k
+ Metrics:
+ mIoU: 49.62
+ mIoU(ms+flip): 50.36
+ Task: Semantic Segmentation
+ Weights: https://download.openmmlab.com/mmsegmentation/v0.5/segformer/segformer_mit-b5_640x640_160k_ade20k/segformer_mit-b5_640x640_160k_ade20k_20210801_121243-41d2845b.pth
diff --git a/docs/useful_tools.md b/docs/useful_tools.md
@@ -255,6 +255,36 @@ Examples:
  python tools/analyze_logs.py log.json --keys loss --legend loss
  ```
 
+### Model conversion
+
+`tools/model_converters/` provide several scripts to convert pretrain models released by other repos to MMSegmentation style.
+
+#### ViT Swin MiT Transformer Models
+
+- ViT
+
+ `tools/model_converters/vit2mmseg.py` convert keys in timm pretrained vit models to MMSegmentation style.
+
+ ```shell
+ python tools/model_converters/vit2mmseg.py ${SRC} ${DST}
+ ```
+
+- Swin
+
+ `tools/model_converters/swin2mmseg.py` convert keys in official pretrained swin models to MMSegmentation style.
+
+ ```shell
+ python tools/model_converters/swin2mmseg.py ${SRC} ${DST}
+ ```
+
+- SegFormer
+
+ `tools/model_converters/mit2mmseg.py` convert keys in official pretrained mit models to MMSegmentation style.
+
+ ```shell
+ python tools/model_converters/mit2mmseg.py ${SRC} ${DST}
+ ```
+
 ## Model Serving
 
 In order to serve an `MMSegmentation` model with [`TorchServe`](https://pytorch.org/serve/), you can follow the steps:

diff --git a/docs_zh-CN/useful_tools.md b/docs_zh-CN/useful_tools.md
@@ -259,6 +259,36 @@ python tools/analyze_logs.py xxx.log.json [--keys ${KEYS}] [--legend ${LEGEND}]
  python tools/analyze_logs.py log.json --keys loss --legend loss
  ```
 
+### 转换其他仓库的权重
+
+`tools/model_converters/` 提供了若干个预训练权重转换脚本，支持将其他仓库的预训练权重的 key 转换为与 MMSegmentation 相匹配的 key。
+
+#### ViT Swin MiT Transformer 模型
+
+- ViT
+
+`tools/model_converters/vit2mmseg.py` 将 timm 预训练模型转换到 MMSegmentation。
+
+ ```shell
+ python tools/model_converters/vit2mmseg.py ${SRC} ${DST}
+ ```
+
+- Swin
+
+ `tools/model_converters/swin2mmseg.py` 将官方预训练模型转换到 MMSegmentation。
+
+ ```shell
+ python tools/model_converters/swin2mmseg.py ${SRC} ${DST}
+ ```
+
+- SegFormer
+
+ `tools/model_converters/mit2mmseg.py` 将官方预训练模型转换到 MMSegmentation。
+
+ ```shell
+ python tools/model_converters/mit2mmseg.py ${SRC} ${DST}
+ ```
+
 ## 模型服务
 
 为了用 [`TorchServe`](https://pytorch.org/serve/) 服务 `MMSegmentation` 的模型 ， 您可以遵循如下流程:

diff --git a/model-index.yml b/model-index.yml
@@ -23,6 +23,7 @@ Import:
 - configs/psanet/psanet.yml
 - configs/pspnet/pspnet.yml
 - configs/resnest/resnest.yml
+- configs/segformer/segformer.yml
 - configs/sem_fpn/sem_fpn.yml
 - configs/setr/setr.yml
 - configs/swin/swin.yml

diff --git a/tools/model_converters/mit_convert.py b/tools/model_converters/mit_convert.py
@@ -5,7 +5,7 @@
 import torch
 
 
-def mit_convert(ckpt):
+def convert_mit(ckpt):
  new_ckpt = OrderedDict()
  # Process the concat between q linear weights and kv linear weights
  for k, v in ckpt.items():
@@ -73,5 +73,5 @@ def parse_args():
 
  ckpt = torch.load(src_path, map_location='cpu')
 
- ckpt = mit_convert(ckpt)
+ ckpt = convert_mit(ckpt)
  torch.save(ckpt, dst_path)
diff --git a/tools/model_converters/swin2mmseg.py b/tools/model_converters/swin2mmseg.py
@@ -0,0 +1,83 @@
+import argparse
+from collections import OrderedDict
+
+import torch
+
+
+def convert_swin(ckpt):
+ new_ckpt = OrderedDict()
+
+ def correct_unfold_reduction_order(x):
+ out_channel, in_channel = x.shape
+ x = x.reshape(out_channel, 4, in_channel // 4)
+ x = x[:, [0, 2, 1, 3], :].transpose(1,
+ 2).reshape(out_channel, in_channel)
+ return x
+
+ def correct_unfold_norm_order(x):
+ in_channel = x.shape[0]
+ x = x.reshape(4, in_channel // 4)
+ x = x[[0, 2, 1, 3], :].transpose(0, 1).reshape(in_channel)
+ return x
+
+ for k, v in ckpt.items():
+ if k.startswith('head'):
+ continue
+ elif k.startswith('layers'):
+ new_v = v
+ if 'attn.' in k:
+ new_k = k.replace('attn.', 'attn.w_msa.')
+ elif 'mlp.' in k:
+ if 'mlp.fc1.' in k:
+ new_k = k.replace('mlp.fc1.', 'ffn.layers.0.0.')
+ elif 'mlp.fc2.' in k:
+ new_k = k.replace('mlp.fc2.', 'ffn.layers.1.')
+ else:
+ new_k = k.replace('mlp.', 'ffn.')
+ elif 'downsample' in k:
+ new_k = k
+ if 'reduction.' in k:
+ new_v = correct_unfold_reduction_order(v)
+ elif 'norm.' in k:
+ new_v = correct_unfold_norm_order(v)
+ else:
+ new_k = k
+ new_k = new_k.replace('layers', 'stages', 1)
+ elif k.startswith('patch_embed'):
+ new_v = v
+ if 'proj' in k:
+ new_k = k.replace('proj', 'projection')
+ else:
+ new_k = k
+ else:
+ new_v = v
+ new_k = k
+
+ new_ckpt[new_k] = new_v
+
+ return new_ckpt
+
+
+def main():
+ parser = argparse.ArgumentParser(
+ description='Convert keys in official pretrained swin models to'
+ 'MMSegmentation style.')
+ parser.add_argument('src', help='src segmentation model path')
+ # The dst path must be a full path of the new checkpoint.
+ parser.add_argument('dst', help='save path')
+ args = parser.parse_args()
+
+ checkpoint = torch.load(args.src, map_location='cpu')
+ if 'state_dict' in checkpoint:
+ state_dict = checkpoint['state_dict']
+ elif 'model' in checkpoint:
+ state_dict = checkpoint['model']
+ else:
+ state_dict = checkpoint
+ weight = convert_swin(state_dict)
+ with open(args.dst, 'wb') as f:
+ torch.save(weight, f)
+
+
+if __name__ == '__main__':
+ main()