Skip to content

Commit 54435fb

Browse files
Junjun2016Junjun2016
authored
[Feature] Support TIMMBackbone (open-mmlab#998)
* add TIMMBackbone and unittests * add timm to tests requirements * deprecate pt1.3.1 * reduce the unittests input of timm backbone * fix ci * fix ci * fix ci * fix ci * fix ci * fix ci * fix ci * fix ci * fix ci * remove unittests of large models of timm backbone * generate coverage report for all unittests env * reduce the unittests input of timm backbone * reduce the unittests input of timm backbone
1 parent ddce375 commit 54435fb

File tree

4 files changed

+215
-1
lines changed

4 files changed

+215
-1
lines changed

‎.github/workflows/build.yml‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,9 +71,17 @@ jobs:
7171
run: rm -rf .eggs && pip install -e .
7272
- name: Run unittests and generate coverage report
7373
run: |
74+
pip install timm
7475
coverage run --branch --source mmseg -m pytest tests/
7576
coverage xml
7677
coverage report -m
78+
if: ${{matrix.torch >= '1.5.0'}}
79+
- name: Skip timm unittests and generate coverage report
80+
run: |
81+
coverage run --branch --source mmseg -m pytest tests/ --ignore tests/test_models/test_backbones/test_timm_backbone.py
82+
coverage xml
83+
coverage report -m
84+
if: ${{matrix.torch < '1.5.0'}}
7785

7886
build_cuda101:
7987
runs-on: ubuntu-18.04
@@ -142,9 +150,17 @@ jobs:
142150
TORCH_CUDA_ARCH_LIST=7.0 pip install .
143151
- name: Run unittests and generate coverage report
144152
run: |
153+
python -m pip install timm
145154
coverage run --branch --source mmseg -m pytest tests/
146155
coverage xml
147156
coverage report -m
157+
if: ${{matrix.torch >= '1.5.0'}}
158+
- name: Skip timm unittests and generate coverage report
159+
run: |
160+
coverage run --branch --source mmseg -m pytest tests/ --ignore tests/test_models/test_backbones/test_timm_backbone.py
161+
coverage xml
162+
coverage report -m
163+
if: ${{matrix.torch < '1.5.0'}}
148164
- name: Upload coverage to Codecov
149165
uses: codecov/[email protected]
150166
with:
@@ -198,6 +214,7 @@ jobs:
198214
TORCH_CUDA_ARCH_LIST=7.0 pip install .
199215
- name: Run unittests and generate coverage report
200216
run: |
217+
python -m pip install timm
201218
coverage run --branch --source mmseg -m pytest tests/
202219
coverage xml
203220
coverage report -m

‎mmseg/models/backbones/__init__.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,13 @@
1212
from .resnet import ResNet, ResNetV1c, ResNetV1d
1313
from .resnext import ResNeXt
1414
from .swin import SwinTransformer
15+
from .timm_backbone import TIMMBackbone
1516
from .unet import UNet
1617
from .vit import VisionTransformer
1718

1819
__all__ = [
1920
'ResNet', 'ResNetV1c', 'ResNetV1d', 'ResNeXt', 'HRNet', 'FastSCNN',
2021
'ResNeSt', 'MobileNetV2', 'UNet', 'CGNet', 'MobileNetV3',
2122
'VisionTransformer', 'SwinTransformer', 'MixVisionTransformer',
22-
'BiSeNetV1', 'BiSeNetV2', 'ICNet'
23+
'BiSeNetV1', 'BiSeNetV2', 'ICNet', 'TIMMBackbone'
2324
]

‎mmseg/models/backbones/timm_backbone.py‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
# Copyright (c) OpenMMLab. All rights reserved.
2+
try:
3+
import timm
4+
except ImportError:
5+
timm = None
6+
7+
from mmcv.cnn.bricks.registry import NORM_LAYERS
8+
from mmcv.runner import BaseModule
9+
10+
from ..builder import BACKBONES
11+
12+
13+
@BACKBONES.register_module()
14+
class TIMMBackbone(BaseModule):
15+
"""Wrapper to use backbones from timm library. More details can be found in
16+
`timm <https://github.com/rwightman/pytorch-image-models>`_ .
17+
18+
Args:
19+
model_name (str): Name of timm model to instantiate.
20+
pretrained (bool): Load pretrained weights if True.
21+
checkpoint_path (str): Path of checkpoint to load after
22+
model is initialized.
23+
in_channels (int): Number of input image channels. Default: 3.
24+
init_cfg (dict, optional): Initialization config dict
25+
**kwargs: Other timm & model specific arguments.
26+
"""
27+
28+
def __init__(
29+
self,
30+
model_name,
31+
features_only=True,
32+
pretrained=True,
33+
checkpoint_path='',
34+
in_channels=3,
35+
init_cfg=None,
36+
**kwargs,
37+
):
38+
if timm is None:
39+
raise RuntimeError('timm is not installed')
40+
super(TIMMBackbone, self).__init__(init_cfg)
41+
if 'norm_layer' in kwargs:
42+
kwargs['norm_layer'] = NORM_LAYERS.get(kwargs['norm_layer'])
43+
self.timm_model = timm.create_model(
44+
model_name=model_name,
45+
features_only=features_only,
46+
pretrained=pretrained,
47+
in_chans=in_channels,
48+
checkpoint_path=checkpoint_path,
49+
**kwargs,
50+
)
51+
52+
# Make unused parameters None
53+
self.timm_model.global_pool = None
54+
self.timm_model.fc = None
55+
self.timm_model.classifier = None
56+
57+
# Hack to use pretrained weights from timm
58+
if pretrained or checkpoint_path:
59+
self._is_init = True
60+
61+
def forward(self, x):
62+
features = self.timm_model(x)
63+
return features

‎tests/test_models/test_backbones/test_timm_backbone.py‎

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
# Copyright (c) OpenMMLab. All rights reserved.
2+
import pytest
3+
import torch
4+
5+
from mmseg.models.backbones import TIMMBackbone
6+
from .utils import check_norm_state
7+
8+
9+
def test_timm_backbone():
10+
with pytest.raises(TypeError):
11+
# pretrained must be a string path
12+
model = TIMMBackbone()
13+
model.init_weights(pretrained=0)
14+
15+
# Test different norm_layer, can be: 'SyncBN', 'BN2d', 'GN', 'LN', 'IN'
16+
# Test resnet18 from timm, norm_layer='BN2d'
17+
model = TIMMBackbone(
18+
model_name='resnet18',
19+
features_only=True,
20+
pretrained=False,
21+
output_stride=32,
22+
norm_layer='BN2d')
23+
24+
# Test resnet18 from timm, norm_layer='SyncBN'
25+
model = TIMMBackbone(
26+
model_name='resnet18',
27+
features_only=True,
28+
pretrained=False,
29+
output_stride=32,
30+
norm_layer='SyncBN')
31+
32+
# Test resnet18 from timm, features_only=True, output_stride=32
33+
model = TIMMBackbone(
34+
model_name='resnet18',
35+
features_only=True,
36+
pretrained=False,
37+
output_stride=32)
38+
model.init_weights()
39+
model.train()
40+
assert check_norm_state(model.modules(), True)
41+
42+
imgs = torch.randn(1, 3, 224, 224)
43+
feats = model(imgs)
44+
feats = [feat.shape for feat in feats]
45+
assert len(feats) == 5
46+
assert feats[0] == torch.Size((1, 64, 112, 112))
47+
assert feats[1] == torch.Size((1, 64, 56, 56))
48+
assert feats[2] == torch.Size((1, 128, 28, 28))
49+
assert feats[3] == torch.Size((1, 256, 14, 14))
50+
assert feats[4] == torch.Size((1, 512, 7, 7))
51+
52+
# Test resnet18 from timm, features_only=True, output_stride=16
53+
model = TIMMBackbone(
54+
model_name='resnet18',
55+
features_only=True,
56+
pretrained=False,
57+
output_stride=16)
58+
imgs = torch.randn(1, 3, 224, 224)
59+
feats = model(imgs)
60+
feats = [feat.shape for feat in feats]
61+
assert len(feats) == 5
62+
assert feats[0] == torch.Size((1, 64, 112, 112))
63+
assert feats[1] == torch.Size((1, 64, 56, 56))
64+
assert feats[2] == torch.Size((1, 128, 28, 28))
65+
assert feats[3] == torch.Size((1, 256, 14, 14))
66+
assert feats[4] == torch.Size((1, 512, 14, 14))
67+
68+
# Test resnet18 from timm, features_only=True, output_stride=8
69+
model = TIMMBackbone(
70+
model_name='resnet18',
71+
features_only=True,
72+
pretrained=False,
73+
output_stride=8)
74+
imgs = torch.randn(1, 3, 224, 224)
75+
feats = model(imgs)
76+
feats = [feat.shape for feat in feats]
77+
assert len(feats) == 5
78+
assert feats[0] == torch.Size((1, 64, 112, 112))
79+
assert feats[1] == torch.Size((1, 64, 56, 56))
80+
assert feats[2] == torch.Size((1, 128, 28, 28))
81+
assert feats[3] == torch.Size((1, 256, 28, 28))
82+
assert feats[4] == torch.Size((1, 512, 28, 28))
83+
84+
# Test efficientnet_b1 with pretrained weights
85+
model = TIMMBackbone(model_name='efficientnet_b1', pretrained=True)
86+
87+
# Test resnetv2_50x1_bitm from timm, features_only=True, output_stride=8
88+
model = TIMMBackbone(
89+
model_name='resnetv2_50x1_bitm',
90+
features_only=True,
91+
pretrained=False,
92+
output_stride=8)
93+
imgs = torch.randn(1, 3, 8, 8)
94+
feats = model(imgs)
95+
feats = [feat.shape for feat in feats]
96+
assert len(feats) == 5
97+
assert feats[0] == torch.Size((1, 64, 4, 4))
98+
assert feats[1] == torch.Size((1, 256, 2, 2))
99+
assert feats[2] == torch.Size((1, 512, 1, 1))
100+
assert feats[3] == torch.Size((1, 1024, 1, 1))
101+
assert feats[4] == torch.Size((1, 2048, 1, 1))
102+
103+
# Test resnetv2_50x3_bitm from timm, features_only=True, output_stride=8
104+
model = TIMMBackbone(
105+
model_name='resnetv2_50x3_bitm',
106+
features_only=True,
107+
pretrained=False,
108+
output_stride=8)
109+
imgs = torch.randn(1, 3, 8, 8)
110+
feats = model(imgs)
111+
feats = [feat.shape for feat in feats]
112+
assert len(feats) == 5
113+
assert feats[0] == torch.Size((1, 192, 4, 4))
114+
assert feats[1] == torch.Size((1, 768, 2, 2))
115+
assert feats[2] == torch.Size((1, 1536, 1, 1))
116+
assert feats[3] == torch.Size((1, 3072, 1, 1))
117+
assert feats[4] == torch.Size((1, 6144, 1, 1))
118+
119+
# Test resnetv2_101x1_bitm from timm, features_only=True, output_stride=8
120+
model = TIMMBackbone(
121+
model_name='resnetv2_101x1_bitm',
122+
features_only=True,
123+
pretrained=False,
124+
output_stride=8)
125+
imgs = torch.randn(1, 3, 8, 8)
126+
feats = model(imgs)
127+
feats = [feat.shape for feat in feats]
128+
assert len(feats) == 5
129+
assert feats[0] == torch.Size((1, 64, 4, 4))
130+
assert feats[1] == torch.Size((1, 256, 2, 2))
131+
assert feats[2] == torch.Size((1, 512, 1, 1))
132+
assert feats[3] == torch.Size((1, 1024, 1, 1))
133+
assert feats[4] == torch.Size((1, 2048, 1, 1))

0 commit comments

Comments
 (0)