gheinrich commited on
Commit
38c3084
1 Parent(s): 3e539bc
adaptor_base.py ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
2
+ #
3
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
4
+ # and proprietary rights in and to this software, related documentation
5
+ # and any modifications thereto. Any use, reproduction, disclosure or
6
+ # distribution of this software and related documentation without an express
7
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
8
+ from argparse import Namespace
9
+ from typing import NamedTuple
10
+
11
+ import torch
12
+ from torch import nn
13
+ import torch.nn.functional as F
14
+
15
+
16
+ class AdaptorInput(NamedTuple):
17
+ images: torch.Tensor
18
+ summary: torch.Tensor
19
+ features: torch.Tensor
20
+
21
+
22
+ class RadioOutput(NamedTuple):
23
+ summary: torch.Tensor
24
+ features: torch.Tensor
25
+
26
+ def to(self, *args, **kwargs):
27
+ return RadioOutput(
28
+ self.summary.to(*args, **kwargs) if self.summary is not None else None,
29
+ self.features.to(*args, **kwargs) if self.features is not None else None,
30
+ )
31
+
32
+
33
+ class AdaptorBase(nn.Module):
34
+ def forward(self, input: AdaptorInput) -> RadioOutput:
35
+ raise NotImplementedError("Subclasses must implement this!")
adaptor_generic.py ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
2
+ #
3
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
4
+ # and proprietary rights in and to this software, related documentation
5
+ # and any modifications thereto. Any use, reproduction, disclosure or
6
+ # distribution of this software and related documentation without an express
7
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
8
+ from argparse import Namespace
9
+
10
+ import torch
11
+ from torch import nn
12
+ import torch.nn.functional as F
13
+
14
+ from .adaptor_base import AdaptorBase, AdaptorInput, RadioOutput
15
+ from .adaptor_mlp import create_mlp_from_state
16
+
17
+
18
+ class GenericAdaptor(AdaptorBase):
19
+ def __init__(self, main_config: Namespace, adaptor_config, state):
20
+ super().__init__()
21
+
22
+ self.head_mlp = create_mlp_from_state(main_config.mlp_version, state, 'summary.')
23
+ self.feat_mlp = create_mlp_from_state(main_config.mlp_version, state, 'feature.')
24
+
25
+ def forward(self, input: AdaptorInput) -> RadioOutput:
26
+ summary = self.head_mlp(input.summary)
27
+ feat = self.feat_mlp(input.features)
28
+
29
+ return RadioOutput(summary, feat)
adaptor_mlp.py ADDED
@@ -0,0 +1,150 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
2
+ #
3
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
4
+ # and proprietary rights in and to this software, related documentation
5
+ # and any modifications thereto. Any use, reproduction, disclosure or
6
+ # distribution of this software and related documentation without an express
7
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
8
+ import math
9
+ from typing import Dict
10
+
11
+ import torch
12
+ from torch import nn
13
+
14
+ from einops import rearrange
15
+ from timm.models.vision_transformer import Block
16
+
17
+
18
+ class MLP(nn.Module):
19
+ def __init__(self, input_size: int, hidden_size: int, output_size: int,
20
+ num_inner: int = 0, device: torch.device = None, **kwargs):
21
+ super(MLP, self).__init__()
22
+ self.fc1 = nn.Linear(input_size, hidden_size, device=device)
23
+ self.norm = nn.LayerNorm(hidden_size, device=device)
24
+ self.relu = nn.ReLU()
25
+
26
+ inner = []
27
+ for _ in range(num_inner):
28
+ inner.extend([
29
+ nn.Linear(hidden_size, hidden_size, device=device),
30
+ nn.LayerNorm(hidden_size, device=device),
31
+ nn.ReLU(),
32
+ ])
33
+ if inner:
34
+ self.inner = nn.Sequential(*inner)
35
+ else:
36
+ self.inner = nn.Identity()
37
+
38
+ self.fc2 = nn.Linear(hidden_size, output_size, device=device)
39
+
40
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
41
+ x = self.fc1(x)
42
+ x = self.norm(x)
43
+ x = self.relu(x)
44
+ x = self.inner(x)
45
+ x = self.fc2(x)
46
+ return x
47
+
48
+
49
+ class MLP2(nn.Module):
50
+ def __init__(self, input_size: int, hidden_size: int, output_size: int,
51
+ num_inner: int = 0,
52
+ pre_norm: bool = False, device: torch.device = None,
53
+ upsample_factor: int = 1,
54
+ **kwargs):
55
+ super().__init__()
56
+
57
+ self.pre_norm = nn.Sequential(
58
+ nn.LayerNorm(input_size),
59
+ nn.GELU(),
60
+ ) if pre_norm else nn.Identity()
61
+
62
+ self.upsample_factor = upsample_factor
63
+ self._real_output_dim = output_size
64
+
65
+ hidden_size *= upsample_factor
66
+ output_size *= (upsample_factor ** 2)
67
+
68
+ self.fc1 = nn.Linear(input_size, hidden_size, device=device)
69
+
70
+ blocks = []
71
+ for _ in range(num_inner):
72
+ blocks.append(nn.Sequential(
73
+ nn.LayerNorm(hidden_size, device=device),
74
+ nn.GELU(),
75
+ nn.Linear(hidden_size, hidden_size, device=device),
76
+ ))
77
+ self.blocks = nn.ModuleList(blocks)
78
+
79
+ self.final = nn.Sequential(
80
+ nn.LayerNorm(hidden_size, device=device),
81
+ nn.GELU(),
82
+ nn.Linear(hidden_size, output_size, device=device),
83
+ )
84
+
85
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
86
+ x = self.pre_norm(x)
87
+ x = self.fc1(x)
88
+ for block in self.blocks:
89
+ x = x + block(x)
90
+ x = self.final(x)
91
+
92
+ if self.upsample_factor > 1:
93
+ h = w = int(math.sqrt(x.shape[1]))
94
+ x = rearrange(x, 'b (h w) (u1 u2 c) -> b (u1 h u2 w) c',
95
+ h=h, w=w, u1=self.upsample_factor, u2=self.upsample_factor,
96
+ c=self._real_output_dim)
97
+
98
+ return x
99
+
100
+
101
+ MLP_FACTORY = {
102
+ 'v1': MLP,
103
+ 'v2': MLP2,
104
+ }
105
+
106
+
107
+ def strip_prefix(state: Dict[str, torch.Tensor], prefix: str):
108
+ state = {
109
+ k[len(prefix):]: v
110
+ for k, v in state.items()
111
+ if k.startswith(prefix)
112
+ }
113
+ return state
114
+
115
+
116
+ def get_mlp_info_from_state(version: str, state: Dict[str, torch.Tensor], prefix: str = ''):
117
+ state = strip_prefix(state, prefix)
118
+
119
+ if version == 'v1':
120
+ hidden_dim, input_dim = state['fc1.weight'].shape
121
+ output_dim = state['fc2.weight'].shape[0]
122
+
123
+ for num_inner in range(1000):
124
+ k = f'inner.{num_inner}.0.weight'
125
+ if k not in state:
126
+ break
127
+ elif version == 'v2':
128
+ hidden_dim, input_dim = state['fc1.weight'].shape
129
+ output_dim = state['final.2.weight'].shape[0]
130
+
131
+ for num_inner in range(1000):
132
+ k = f'blocks.{num_inner}.0.weight'
133
+ if k not in state:
134
+ break
135
+ else:
136
+ raise ValueError(f'Unsupported MLP version: {version}')
137
+
138
+ return input_dim, hidden_dim, output_dim, num_inner
139
+
140
+
141
+ def create_mlp_from_state(version: str, state: Dict[str, torch.Tensor], prefix: str = ''):
142
+ state = strip_prefix(state, prefix)
143
+
144
+ input_dim, hidden_dim, output_dim, num_inner = get_mlp_info_from_state(version, state)
145
+
146
+ ret: nn.Module = MLP_FACTORY[version](input_dim, hidden_dim, output_dim, num_inner)
147
+
148
+ ret.load_state_dict(state)
149
+
150
+ return ret
adaptor_registry.py ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
2
+ #
3
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
4
+ # and proprietary rights in and to this software, related documentation
5
+ # and any modifications thereto. Any use, reproduction, disclosure or
6
+ # distribution of this software and related documentation without an express
7
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
8
+ from argparse import Namespace
9
+ from typing import Dict, Any
10
+
11
+ import torch
12
+
13
+ from .adaptor_generic import GenericAdaptor, AdaptorBase
14
+
15
+ dict_t = Dict[str, Any]
16
+ state_t = Dict[str, torch.Tensor]
17
+
18
+
19
+ class AdaptorRegistry:
20
+ def __init__(self):
21
+ self._registry = {}
22
+
23
+ def register_adaptor(self, name):
24
+ def decorator(factory_function):
25
+ if name in self._registry:
26
+ raise ValueError(f"Model '{name}' already registered")
27
+ self._registry[name] = factory_function
28
+ return factory_function
29
+ return decorator
30
+
31
+ def create_adaptor(self, name, main_config: Namespace, adaptor_config: dict_t, state: state_t) -> AdaptorBase:
32
+ if name not in self._registry:
33
+ return GenericAdaptor(main_config, adaptor_config, state)
34
+ return self._registry[name](main_config, adaptor_config, state)
35
+
36
+ # Creating an instance of the registry
37
+ adaptor_registry = AdaptorRegistry()
cls_token.py ADDED
@@ -0,0 +1,55 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved.
2
+ #
3
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
4
+ # and proprietary rights in and to this software, related documentation
5
+ # and any modifications thereto. Any use, reproduction, disclosure or
6
+ # distribution of this software and related documentation without an express
7
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
8
+
9
+ import torch
10
+ from torch import nn
11
+
12
+
13
+ class ClsToken(nn.Module):
14
+ def __init__(self, ndim: int,
15
+ num_tokens: int = 1,
16
+ enabled: bool = True,
17
+ register_multiple: int = 0,
18
+ ):
19
+ super().__init__()
20
+
21
+ self.ndim = ndim
22
+ self.enabled = enabled
23
+ self.num_registers = 0
24
+ self.num_tokens = num_tokens
25
+ if enabled:
26
+ if register_multiple > 0:
27
+ self.num_registers = register_multiple - (num_tokens % register_multiple)
28
+
29
+ scale = ndim ** -0.5
30
+ self.token = nn.Parameter(torch.randn(num_tokens + self.num_registers, ndim) * scale)
31
+ else:
32
+ self.token = None
33
+
34
+ self.num_patches = self.num_tokens + self.num_registers
35
+
36
+ def disable(self):
37
+ self.token = None
38
+ self.enabled = False
39
+
40
+ def forward(self, x: torch.Tensor):
41
+ if self.token is None:
42
+ return x
43
+
44
+ token = self.token.unsqueeze(0).expand(x.shape[0], -1, -1)
45
+ x = torch.cat([
46
+ token,
47
+ x,
48
+ ], dim=1)
49
+
50
+ return x
51
+
52
+ def no_weight_decay(self):
53
+ return [
54
+ 'token',
55
+ ]
common.py ADDED
@@ -0,0 +1,71 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
2
+ #
3
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
4
+ # and proprietary rights in and to this software, related documentation
5
+ # and any modifications thereto. Any use, reproduction, disclosure or
6
+ # distribution of this software and related documentation without an express
7
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
8
+
9
+ from dataclasses import dataclass
10
+ from typing import Optional
11
+
12
+ from .radio_model import Resolution
13
+
14
+
15
+ @dataclass
16
+ class RadioResource:
17
+ url: str
18
+ patch_size: int
19
+ max_resolution: int
20
+ preferred_resolution: Resolution
21
+ vitdet_num_windowed: Optional[int] = None
22
+ vitdet_num_global: Optional[int] = None
23
+
24
+
25
+ RESOURCE_MAP = {
26
+ # RADIOv2.5
27
+ "radio_v2.5-b": RadioResource(
28
+ "https://huggingface.co/nvidia/RADIO/resolve/main/radio-v2.5-b_half.pth.tar?download=true",
29
+ patch_size=16,
30
+ max_resolution=2048,
31
+ preferred_resolution=(768, 768),
32
+ vitdet_num_global=4,
33
+ ),
34
+ "radio_v2.5-l": RadioResource(
35
+ "https://huggingface.co/nvidia/RADIO/resolve/main/radio-v2.5-l_half.pth.tar?download=true",
36
+ patch_size=16,
37
+ max_resolution=2048,
38
+ preferred_resolution=(768, 768),
39
+ vitdet_num_global=4,
40
+ ),
41
+ # RADIO
42
+ "radio_v2.1": RadioResource(
43
+ "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.1_bf16.pth.tar?download=true",
44
+ patch_size=16,
45
+ max_resolution=2048,
46
+ preferred_resolution=Resolution(432, 432),
47
+ vitdet_num_windowed=5,
48
+ ),
49
+ "radio_v2": RadioResource(
50
+ "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.pth.tar?download=true",
51
+ patch_size=16,
52
+ max_resolution=2048,
53
+ preferred_resolution=Resolution(432, 432),
54
+ vitdet_num_windowed=5,
55
+ ),
56
+ "radio_v1": RadioResource(
57
+ "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v1.pth.tar?download=true",
58
+ patch_size=14,
59
+ max_resolution=1050,
60
+ preferred_resolution=Resolution(378, 378),
61
+ ),
62
+ # E-RADIO
63
+ "e-radio_v2": RadioResource(
64
+ "https://huggingface.co/nvidia/RADIO/resolve/main/eradio_v2.pth.tar?download=true",
65
+ patch_size=16,
66
+ max_resolution=2048,
67
+ preferred_resolution=Resolution(512, 512),
68
+ ),
69
+ }
70
+
71
+ DEFAULT_VERSION = "radio_v2.5-l"
config.json ADDED
@@ -0,0 +1,221 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "adaptor_names": null,
3
+ "architectures": [
4
+ "RADIOModel"
5
+ ],
6
+ "args": {
7
+ "aa": null,
8
+ "amp": true,
9
+ "amp_dtype": "bfloat16",
10
+ "amp_impl": "native",
11
+ "aug_repeats": 0,
12
+ "aug_splits": 0,
13
+ "bn_eps": null,
14
+ "bn_momentum": null,
15
+ "cache_dir": null,
16
+ "channels_last": false,
17
+ "checkpoint_hist": 10,
18
+ "chk_keep_forever": 50,
19
+ "class_map": "",
20
+ "clip_grad": null,
21
+ "clip_mode": "norm",
22
+ "cls_token_per_teacher": true,
23
+ "coco_annotations_file": "/datasets/coco2017-adlsa/annotations/captions_val2017.json",
24
+ "coco_image_dir": "/datasets/coco2017-adlsa/val2017",
25
+ "color_jitter": 0.4,
26
+ "cooldown_epochs": 0,
27
+ "cpe_max_size": 2048,
28
+ "crd_loss": false,
29
+ "crd_loss_weight": 0.8,
30
+ "crop_pct": null,
31
+ "cutmix": 0.0,
32
+ "cutmix_minmax": null,
33
+ "dataset_download": false,
34
+ "debug_full_knn": false,
35
+ "decay_epochs": 90,
36
+ "decay_milestones": [
37
+ 90,
38
+ 180,
39
+ 270
40
+ ],
41
+ "decay_rate": 0.1,
42
+ "device": "cuda:0",
43
+ "dist_bn": "reduce",
44
+ "dist_norm_weight": 0.0,
45
+ "distributed": true,
46
+ "drop": 0.0,
47
+ "drop_block": null,
48
+ "drop_connect": null,
49
+ "drop_path": null,
50
+ "dtype": "float32",
51
+ "epoch_repeats": 0.0,
52
+ "eval": false,
53
+ "eval_metric": "knn_top1",
54
+ "eval_teacher": false,
55
+ "eval_teacher_only": false,
56
+ "eval_throughput": false,
57
+ "fast_norm": false,
58
+ "fd_loss_fn": "MSE",
59
+ "feature_normalization": "SHIP_NORM",
60
+ "feature_summarizer": "cls_token",
61
+ "feature_upscale_factor": null,
62
+ "force_new_wandb_id": false,
63
+ "force_spectral_reparam": true,
64
+ "freeze_bn": false,
65
+ "fsdp": false,
66
+ "fuser": "",
67
+ "gp": null,
68
+ "grad_accum_steps": 1,
69
+ "grad_checkpointing": false,
70
+ "head_init_bias": null,
71
+ "head_init_scale": null,
72
+ "head_warmup": 5,
73
+ "head_weight_decay": 5e-05,
74
+ "hflip": 0.5,
75
+ "img_size": null,
76
+ "in_chans": null,
77
+ "initial_checkpoint": "",
78
+ "input_size": null,
79
+ "interpolation": "",
80
+ "layer_decay": null,
81
+ "local_rank": 0,
82
+ "log_interval": 50,
83
+ "log_mlflow": false,
84
+ "log_wandb": true,
85
+ "loss_auto_balance": false,
86
+ "lr_base": 0.1,
87
+ "lr_base_scale": "",
88
+ "lr_base_size": 256,
89
+ "lr_cycle_decay": 0.5,
90
+ "lr_cycle_limit": 1,
91
+ "lr_cycle_mul": 1.0,
92
+ "lr_k_decay": 1.0,
93
+ "lr_noise": null,
94
+ "lr_noise_pct": 0.67,
95
+ "lr_noise_std": 1.0,
96
+ "mean": null,
97
+ "mesa": false,
98
+ "min_lr": 0,
99
+ "mixup": 0.0,
100
+ "mixup_mode": "batch",
101
+ "mixup_off_epoch": 0,
102
+ "mixup_prob": 1.0,
103
+ "mixup_switch_prob": 0.5,
104
+ "mlp_hidden_size": 1520,
105
+ "mlp_num_inner": 1,
106
+ "mlp_version": "v2",
107
+ "model": "vit_base_patch16_224",
108
+ "model_kwargs": {},
109
+ "model_norm": false,
110
+ "momentum": 0.9,
111
+ "no_aug": false,
112
+ "no_ddp_bb": true,
113
+ "no_prefetcher": false,
114
+ "no_resume_opt": false,
115
+ "num_classes": null,
116
+ "opt_betas": null,
117
+ "opt_eps": null,
118
+ "patience_epochs": 10,
119
+ "pin_mem": false,
120
+ "prefetcher": true,
121
+ "pretrained": false,
122
+ "rank": 0,
123
+ "ratio": [
124
+ 0.75,
125
+ 1.3333333333333333
126
+ ],
127
+ "recount": 1,
128
+ "recovery_interval": 0,
129
+ "register_multiple": 8,
130
+ "remode": "pixel",
131
+ "reprob": 0.0,
132
+ "reset_loss_state": true,
133
+ "resplit": false,
134
+ "save_images": false,
135
+ "scale": [
136
+ 0.5,
137
+ 1.0
138
+ ],
139
+ "sched": "cosine",
140
+ "seed": 42,
141
+ "smoothing": 0.1,
142
+ "spectral_reparam": false,
143
+ "split_bn": false,
144
+ "start_epoch": null,
145
+ "std": null,
146
+ "sync_bn": false,
147
+ "synchronize_step": false,
148
+ "teachers": [
149
+ {
150
+ "fd_normalize": false,
151
+ "feature_distillation": true,
152
+ "input_size": 378,
153
+ "model": "ViT-H-14-378-quickgelu",
154
+ "name": "clip",
155
+ "pretrained": "dfn5b",
156
+ "type": "open_clip"
157
+ },
158
+ {
159
+ "fd_normalize": false,
160
+ "feature_distillation": true,
161
+ "input_size": 378,
162
+ "model": "ViT-SO400M-14-SigLIP-384",
163
+ "name": "siglip",
164
+ "pretrained": "webli",
165
+ "type": "open_clip",
166
+ "use_summary": true
167
+ },
168
+ {
169
+ "fd_normalize": false,
170
+ "feature_distillation": true,
171
+ "input_size": 378,
172
+ "model": "dinov2_vitg14_reg",
173
+ "name": "dino_v2",
174
+ "type": "dino_v2"
175
+ },
176
+ {
177
+ "fd_normalize": false,
178
+ "feature_distillation": true,
179
+ "input_size": 1024,
180
+ "model": "vit-h",
181
+ "name": "sam",
182
+ "type": "sam",
183
+ "use_summary": false
184
+ }
185
+ ],
186
+ "torchcompile": null,
187
+ "torchscript": false,
188
+ "train_interpolation": "random",
189
+ "train_split": "train",
190
+ "tta": 0,
191
+ "use_coco": false,
192
+ "use_multi_epochs_loader": false,
193
+ "val_ema_only": false,
194
+ "val_split": "val",
195
+ "vflip": 0.0,
196
+ "vitdet_version": 1,
197
+ "wandb_entity": "",
198
+ "wandb_job_type": "",
199
+ "wandb_name": "",
200
+ "wandb_project": "",
201
+ "warmup_lr": 1e-05,
202
+ "warmup_prefix": false,
203
+ "worker_seeding": "all",
204
+ "workers": 8,
205
+ "world_size": 128
206
+ },
207
+ "auto_map": {
208
+ "AutoConfig": "hf_model.RADIOConfig",
209
+ "AutoModel": "hf_model.RADIOModel"
210
+ },
211
+ "max_resolution": 2048,
212
+ "patch_size": 16,
213
+ "preferred_resolution": [
214
+ 768,
215
+ 768
216
+ ],
217
+ "torch_dtype": "float32",
218
+ "transformers_version": "4.40.1",
219
+ "version": "radio_v2.5-b",
220
+ "vitdet_window_size": null
221
+ }
enable_cpe_support.py ADDED
@@ -0,0 +1,158 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved.
2
+ #
3
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
4
+ # and proprietary rights in and to this software, related documentation
5
+ # and any modifications thereto. Any use, reproduction, disclosure or
6
+ # distribution of this software and related documentation without an express
7
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
8
+
9
+ from typing import List, Optional, Set, Tuple, Union
10
+ from types import MethodType
11
+
12
+ import torch
13
+ from torch import nn
14
+
15
+ from timm.models import VisionTransformer, checkpoint_seq
16
+
17
+ from .vit_patch_generator import ViTPatchGenerator
18
+
19
+
20
+ def _forward_cpe(self: VisionTransformer, x: torch.Tensor) -> torch.Tensor:
21
+ x = self.patch_generator(x)
22
+ if self.grad_checkpointing and not torch.jit.is_scripting():
23
+ x = checkpoint_seq(self.blocks, x)
24
+ else:
25
+ x = self.blocks(x)
26
+ x = self.norm(x)
27
+ return x
28
+
29
+
30
+ def _take_indices(
31
+ num_blocks: int,
32
+ n: Optional[Union[int, List[int], Tuple[int]]],
33
+ ) -> Tuple[Set[int], int]:
34
+ if isinstance(n, int):
35
+ assert n >= 0
36
+ take_indices = {x for x in range(num_blocks - n, num_blocks)}
37
+ else:
38
+ take_indices = {num_blocks + idx if idx < 0 else idx for idx in n}
39
+ return take_indices, max(take_indices)
40
+
41
+
42
+ def _forward_intermediates_cpe(
43
+ self,
44
+ x: torch.Tensor,
45
+ indices: Optional[Union[int, List[int], Tuple[int]]] = None,
46
+ return_prefix_tokens: bool = False,
47
+ norm: bool = False,
48
+ stop_early: bool = False,
49
+ output_fmt: str = 'NCHW',
50
+ intermediates_only: bool = False,
51
+ aggregation: Optional[str] = "sparse",
52
+ ) -> Union[List[torch.Tensor], Tuple[torch.Tensor, List[torch.Tensor]]]:
53
+ """ Forward features that returns intermediates.
54
+
55
+ The Dense layer aggregation method is inspired from the paper: "Dense Connector for MLLMs"
56
+ by Yao, Huanjin et al. (2024). arXiv preprint arXiv:2405.13800}
57
+
58
+ Args:
59
+ x: Input image tensor
60
+ indices: Take last n blocks if int, select matching indices if sequence
61
+ return_prefix_tokens: Return both prefix and spatial intermediate tokens
62
+ norm: Apply norm layer to all intermediates
63
+ stop_early: Stop iterating over blocks when last desired intermediate hit
64
+ output_fmt: Shape of intermediate feature outputs
65
+ intermediates_only: Only return intermediate features
66
+ aggregation: intermediate layer aggregation method (sparse or dense)
67
+ Returns:
68
+ """
69
+ assert output_fmt in ('NCHW', 'NLC'), 'Output format must be one of NCHW or NLC.'
70
+ assert aggregation in ('sparse', 'dense'), 'Aggregation must be one of sparse or dense.'
71
+ reshape = output_fmt == 'NCHW'
72
+ intermediates = []
73
+ take_indices, max_index = _take_indices(len(self.blocks), indices)
74
+ # forward pass
75
+ B, _, height, width = x.shape
76
+ x = self.patch_generator(x)
77
+
78
+ if not stop_early: # can't slice blocks in torchscript
79
+ blocks = self.blocks
80
+ else:
81
+ blocks = self.blocks[:max_index + 1]
82
+
83
+ accumulator = 0
84
+ num_accumulated = 0
85
+
86
+ for i, blk in enumerate(blocks):
87
+ x = blk(x)
88
+ if aggregation == "dense":
89
+ accumulator = accumulator + x
90
+ num_accumulated += 1
91
+ if i in take_indices:
92
+ if aggregation == "dense":
93
+ x_ = accumulator / num_accumulated
94
+ num_accumulated = 0
95
+ accumulator = 0
96
+ else:
97
+ x_ = x
98
+ # normalize intermediates with final norm layer if enabled
99
+ intermediates.append(self.norm(x_) if norm else x_)
100
+
101
+ # process intermediates
102
+
103
+ # split prefix (e.g. class, distill) and spatial feature tokens
104
+ prefix_tokens = [y[:, 0:self.patch_generator.num_cls_tokens] for y in intermediates]
105
+ intermediates = [y[:, self.patch_generator.num_skip:] for y in intermediates]
106
+
107
+ if reshape:
108
+ # reshape to BCHW output format
109
+ H = height // self.patch_generator.patch_size
110
+ W = width // self.patch_generator.patch_size
111
+ intermediates = [y.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() for y in intermediates]
112
+ if not torch.jit.is_scripting() and return_prefix_tokens:
113
+ # return_prefix not support in torchscript due to poor type handling
114
+ intermediates = list(zip(intermediates, prefix_tokens))
115
+ if intermediates_only:
116
+ return intermediates
117
+ x = self.norm(x)
118
+ return x, intermediates
119
+
120
+ def enable_cpe(model: nn.Module,
121
+ max_img_size: Union[int, Tuple[int, int]] = 1024,
122
+ num_cls_tokens: int = 1,
123
+ pos_dropout: float = 0.1,
124
+ register_multiple: int = 0,
125
+ ):
126
+ if not isinstance(model, VisionTransformer):
127
+ raise ValueError("CPE only support for VisionTransformer models!")
128
+
129
+ patch_size = model.patch_embed.patch_size[0]
130
+ embed_dim = model.embed_dim
131
+ input_dims = model.patch_embed.img_size
132
+ normalize_patches = not isinstance(model.patch_embed.norm, nn.Identity)
133
+ cls_token = model.cls_token is not None
134
+
135
+ max_img_size = int(round(max_img_size / patch_size) * patch_size)
136
+
137
+ patch_generator = ViTPatchGenerator(
138
+ patch_size=patch_size,
139
+ embed_dim=embed_dim,
140
+ input_dims=input_dims,
141
+ normalize_patches=normalize_patches,
142
+ cls_token=cls_token,
143
+ max_input_dims=max_img_size,
144
+ pos_dropout=pos_dropout,
145
+ num_cls_tokens=num_cls_tokens,
146
+ register_multiple=register_multiple,
147
+ )
148
+
149
+ model.patch_generator = patch_generator
150
+ model.patch_embed = None
151
+ model.cls_token = None
152
+ model.pos_embed = None
153
+ model.pos_drop = None
154
+ model.num_cls_tokens = num_cls_tokens
155
+ model.num_registers = patch_generator.num_registers
156
+
157
+ model.forward_features = MethodType(_forward_cpe, model)
158
+ model.forward_intermediates = MethodType(_forward_intermediates_cpe, model)
enable_spectral_reparam.py ADDED
@@ -0,0 +1,227 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from logging import getLogger
2
+ import math
3
+ import os
4
+ from typing import Union, Tuple
5
+ from types import MethodType
6
+
7
+ import torch
8
+ from torch import nn
9
+ from torch.nn import functional as F
10
+ from torch.nn.utils import parametrize
11
+ from torch.nn.utils.parametrizations import _SpectralNorm
12
+
13
+ from timm.models.vision_transformer import Attention, Mlp
14
+
15
+ _EPS = 1e-5
16
+
17
+
18
+ class _SNReweight(_SpectralNorm):
19
+ def __init__(self, weight: torch.Tensor, *args, init_norm_to_current: bool = False, alpha: float = 0.05, version: int = 2, **kwargs):
20
+ super().__init__(weight, *args, **kwargs)
21
+
22
+ self.alpha = alpha
23
+ self.version = version
24
+ self.register_buffer('_sn_version', torch.tensor(version))
25
+
26
+ if init_norm_to_current:
27
+ # This will set the numerator to match the denominator, which should preserve the original values
28
+ init_scale = self._get_sigma(weight).item()
29
+ else:
30
+ init_scale = 1.0
31
+
32
+ if version == 1:
33
+ init_value = init_scale
34
+ elif version == 2:
35
+ t = init_scale - alpha
36
+ if t < _EPS:
37
+ getLogger("spectral_reparam").warn(f'The initialized spectral norm {init_scale} is too small to be represented. Setting to {_EPS} instead.')
38
+ t = _EPS
39
+
40
+ init_value = math.log(math.exp(t) - 1)
41
+ else:
42
+ raise ValueError(f'Unsupported version: {version}')
43
+
44
+ # Make 2D so that weight decay gets applied
45
+ self.scale = nn.Parameter(torch.tensor([[init_value]], dtype=torch.float32, device=weight.device))
46
+
47
+ # Re-implementing this because we need to make division by sigma safe
48
+ def _get_sigma(self, weight: torch.Tensor) -> torch.Tensor:
49
+ if weight.ndim == 1:
50
+ # Faster and more exact path, no need to approximate anything
51
+ sigma = weight.norm()
52
+ else:
53
+ weight_mat = self._reshape_weight_to_matrix(weight)
54
+ if self.training:
55
+ self._power_method(weight_mat, self.n_power_iterations)
56
+ # See above on why we need to clone
57
+ u = self._u.clone(memory_format=torch.contiguous_format)
58
+ v = self._v.clone(memory_format=torch.contiguous_format)
59
+ # The proper way of computing this should be through F.bilinear, but
60
+ # it seems to have some efficiency issues:
61
+ # https://github.com/pytorch/pytorch/issues/58093
62
+ sigma = torch.dot(u, torch.mv(weight_mat, v))
63
+
64
+ return sigma + self.eps
65
+
66
+ def forward(self, weight: torch.Tensor, *args, **kwargs):
67
+ dtype = weight.dtype
68
+ sigma = self._get_sigma(weight, *args, **kwargs)
69
+
70
+ if self.version == 1:
71
+ scale = self.scale
72
+ elif self.version == 2:
73
+ scale = F.softplus(self.scale) + self.alpha
74
+ else:
75
+ raise ValueError(f'Unsupported version: {self.version}')
76
+
77
+ scale = scale.float() / sigma.float()
78
+
79
+ y = weight * scale
80
+
81
+ if dtype in (torch.float16, torch.bfloat16):
82
+ y = y.to(dtype)
83
+ return y
84
+
85
+ def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
86
+ version_key = f'{prefix}_sn_version'
87
+ if version_key not in state_dict:
88
+ self.version = 1
89
+ state_dict[version_key] = torch.tensor(1)
90
+ return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
91
+
92
+
93
+ class _AttnSNReweight(nn.Module):
94
+ def __init__(self, weight: torch.Tensor, *args, init_norm_to_current: bool = False, renorm_values: bool = False, **kwargs):
95
+ super().__init__()
96
+
97
+ parts = weight.split(weight.shape[0] // 3, dim=0)
98
+
99
+ ct = 2 if not renorm_values else 3
100
+
101
+ self.parts = nn.ModuleList([
102
+ _SNReweight(p, *args, init_norm_to_current=init_norm_to_current, **kwargs) if i < ct else nn.Identity()
103
+ for i, p in enumerate(parts)
104
+ ])
105
+
106
+ def forward(self, weight: torch.Tensor, *args, **kwargs):
107
+ parts = weight.split(weight.shape[0] // 3, dim=0)
108
+
109
+ parts = [
110
+ fn(p)
111
+ for fn, p in zip(self.parts, parts)
112
+ ]
113
+
114
+ return torch.cat(parts, dim=0)
115
+
116
+
117
+ def enable_spectral_reparam(model: nn.Module,
118
+ n_power_iterations: int = 1,
119
+ eps: float = 1e-6,
120
+ init_norm_to_current: bool = False,
121
+ renorm_values: bool = True,
122
+ renorm_mlp: bool = True):
123
+ # print('Enabling spectral reparametrization')
124
+ for mod in model.modules():
125
+ if isinstance(mod, Attention):
126
+ parametrize.register_parametrization(
127
+ mod.qkv,
128
+ 'weight',
129
+ _AttnSNReweight(mod.qkv.weight, n_power_iterations, dim=0, eps=eps, init_norm_to_current=init_norm_to_current, renorm_values=renorm_values),
130
+ )
131
+ pass
132
+ elif isinstance(mod, Mlp) and renorm_mlp:
133
+ parametrize.register_parametrization(
134
+ mod.fc1,
135
+ 'weight',
136
+ _SNReweight(mod.fc1.weight, n_power_iterations, dim=0, eps=eps, init_norm_to_current=init_norm_to_current),
137
+ )
138
+ parametrize.register_parametrization(
139
+ mod.fc2,
140
+ 'weight',
141
+ _SNReweight(mod.fc2.weight, n_power_iterations, dim=0, eps=eps, init_norm_to_current=init_norm_to_current),
142
+ )
143
+ pass
144
+
145
+
146
+ def configure_spectral_reparam_from_args(model: nn.Module, args):
147
+ spectral_reparam = getattr(args, 'spectral_reparam', False)
148
+ if isinstance(spectral_reparam, bool) and spectral_reparam:
149
+ enable_spectral_reparam(model, init_norm_to_current=args.pretrained)
150
+ elif isinstance(spectral_reparam, dict):
151
+ enable_spectral_reparam(
152
+ model,
153
+ n_power_iterations=spectral_reparam.get('n_power_iterations', 1),
154
+ eps=spectral_reparam.get('eps', 1e-12),
155
+ init_norm_to_current=args.pretrained,
156
+ )
157
+
158
+
159
+ def disable_spectral_reparam(model: nn.Module):
160
+ for mod in model.modules():
161
+ if isinstance(mod, Attention):
162
+ parametrize.remove_parametrizations(mod.qkv, 'weight')
163
+ pass
164
+ elif isinstance(mod, Mlp):
165
+ parametrize.remove_parametrizations(mod.fc1, 'weight')
166
+ parametrize.remove_parametrizations(mod.fc2, 'weight')
167
+ pass
168
+
169
+
170
+ if __name__ == '__main__':
171
+ import argparse
172
+ from . import radio_model as create_model
173
+
174
+ parser = argparse.ArgumentParser(description='Remove parametrization from state dict')
175
+ parser.add_argument('--checkpoint', type=str, required=True, help='The checkpoint to load')
176
+ parser.add_argument('--output', type=str, default='', help='Where to store the checkpoint')
177
+ parser.add_argument('--release', default=False, action='store_true', help='Prune extraneous checkpoint fields')
178
+ parser.add_argument('--strict', default=False, action='store_true', help='Strictly load the state dict')
179
+
180
+ args = parser.parse_args()
181
+
182
+ if not args.output:
183
+ chk_dir, chk_name = os.path.split(args.checkpoint)
184
+ args.output = os.path.join(chk_dir, f'clean_{chk_name}')
185
+ print(f'Set output to "{args.output}"')
186
+
187
+ chk = torch.load(args.checkpoint, map_location='cpu', mmap=True)
188
+
189
+ model = create_model.create_model_from_args(chk['args'])
190
+
191
+ key = 'base_model.'
192
+ mod_state = dict()
193
+ extra_state = dict()
194
+ for k, v in chk['state_dict'].items():
195
+ if k.startswith(key):
196
+ mod_state[k[len(key):]] = v
197
+ else:
198
+ extra_state[k] = v
199
+
200
+ chk_load_info = model.load_state_dict(mod_state, strict=args.strict)
201
+ if chk_load_info.unexpected_keys or chk_load_info.missing_keys:
202
+ print(chk_load_info)
203
+
204
+ if chk['args'].spectral_reparam:
205
+ disable_spectral_reparam(model)
206
+
207
+ if hasattr(chk['args'], 'dtype'):
208
+ model.to(dtype=chk['args'].dtype)
209
+
210
+ mod_state = model.state_dict()
211
+ final_state = dict()
212
+ final_state.update({f'{key}{k}': v for k, v in mod_state.items()})
213
+ final_state.update(extra_state)
214
+
215
+ chk['state_dict'] = final_state
216
+ chk['args'].spectral_reparam = False
217
+
218
+ if args.release:
219
+ chk = {
220
+ 'arch': chk['arch'],
221
+ 'epoch': chk['epoch'],
222
+ 'state_dict': chk['state_dict'],
223
+ 'args': chk['args'],
224
+ }
225
+
226
+ torch.save(chk, args.output)
227
+ pass
eradio_model.py ADDED
@@ -0,0 +1,1392 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+
3
+ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
4
+ #
5
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
6
+ # and proprietary rights in and to this software, related documentation
7
+ # and any modifications thereto. Any use, reproduction, disclosure or
8
+ # distribution of this software and related documentation without an express
9
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
10
+
11
+ # E-RADIO model from
12
+ # Mike Ranzinger, Greg Heinrich, Jan Kautz, and Pavlo Molchanov. "AM-RADIO: Agglomerative Model--Reduce All Domains Into One." arXiv preprint arXiv:2312.06709 (2023).
13
+
14
+ # based on FasterViT, Swin Transformer, YOLOv8
15
+
16
+ # FasterViT:
17
+ # Ali Hatamizadeh, Greg Heinrich, Hongxu Yin, Andrew Tao, Jose M. Alvarez, Jan Kautz, and Pavlo Molchanov. "FasterViT: Fast Vision Transformers with Hierarchical Attention." arXiv preprint arXiv:2306.06189 (2023).
18
+
19
+ import timm
20
+ import torch
21
+ import torch.nn as nn
22
+ from timm.models.registry import register_model
23
+
24
+ from timm.models.layers import trunc_normal_, DropPath, LayerNorm2d
25
+ import numpy as np
26
+ import torch.nn.functional as F
27
+ import math
28
+ import warnings
29
+
30
+ #######################
31
+ ## Codebase from YOLOv8
32
+ ## BEGINNING
33
+ #######################
34
+
35
+ class C2f(nn.Module):
36
+ """Faster Implementation of CSP Bottleneck with 2 convolutions."""
37
+ """From YOLOv8 codebase"""
38
+ def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5, drop_path=None): # ch_in, ch_out, number, shortcut, groups, expansion
39
+ super().__init__()
40
+ if drop_path is None:
41
+ drop_path = [0.0] * n
42
+
43
+ self.c = int(c2 * e) # hidden channels
44
+ self.cv1 = Conv(c1, 2 * self.c, 1, 1)
45
+ self.cv2 = Conv((2 + n) * self.c, c2, 1) # optional act=FReLU(c2)
46
+ self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0, drop_path=drop_path[i]) for i in range(n))
47
+
48
+ def forward(self, x):
49
+ """Forward pass through C2f layer."""
50
+ y = list(self.cv1(x).chunk(2, 1))
51
+ y.extend(m(y[-1]) for m in self.m)
52
+ return self.cv2(torch.cat(y, 1))
53
+
54
+ def forward_split(self, x):
55
+ """Forward pass using split() instead of chunk()."""
56
+ y = list(self.cv1(x).split((self.c, self.c), 1))
57
+ y.extend(m(y[-1]) for m in self.m)
58
+ return self.cv2(torch.cat(y, 1))
59
+
60
+ class Bottleneck(nn.Module):
61
+ """Standard bottleneck."""
62
+
63
+ def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5, drop_path=0.0): # ch_in, ch_out, shortcut, groups, kernels, expand
64
+ super().__init__()
65
+ c_ = int(c2 * e) # hidden channels
66
+ self.cv1 = Conv(c1, c_, k[0], 1)
67
+ self.cv2 = Conv(c_, c2, k[1], 1, g=g)
68
+ self.add = shortcut and c1 == c2
69
+ self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
70
+
71
+ def forward(self, x):
72
+ """'forward()' applies the YOLOv5 FPN to input data."""
73
+ return x + self.drop_path1(self.cv2(self.cv1(x))) if self.add else self.cv2(self.cv1(x))
74
+
75
+
76
+ class Conv(nn.Module):
77
+ """Modified to support layer fusion"""
78
+ default_act = nn.SiLU() # default activation
79
+
80
+ def __init__(self, a, b, kernel_size=1, stride=1, padding=None, g=1, dilation=1, bn_weight_init=1, bias=False, act=True):
81
+ super().__init__()
82
+
83
+ self.conv = torch.nn.Conv2d(a, b, kernel_size, stride, autopad(kernel_size, padding, dilation), dilation, g, bias=False)
84
+ if 1:
85
+ self.bn = torch.nn.BatchNorm2d(b)
86
+ torch.nn.init.constant_(self.bn.weight, bn_weight_init)
87
+ torch.nn.init.constant_(self.bn.bias, 0)
88
+ self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
89
+
90
+
91
+ def forward(self,x):
92
+ x = self.conv(x)
93
+ x = self.bn(x)
94
+ x = self.act(x)
95
+ return x
96
+
97
+ @torch.no_grad()
98
+ def switch_to_deploy(self):
99
+ # return 1
100
+ if not isinstance(self.bn, nn.Identity):
101
+ c, bn = self.conv, self.bn
102
+ w = bn.weight / (bn.running_var + bn.eps) ** 0.5
103
+ w = c.weight * w[:, None, None, None]
104
+ b = bn.bias - bn.running_mean * bn.weight / \
105
+ (bn.running_var + bn.eps)**0.5
106
+
107
+ self.conv.weight.data.copy_(w)
108
+ self.conv.bias = nn.Parameter(b)
109
+
110
+ self.bn = nn.Identity()
111
+
112
+ def autopad(k, p=None, d=1): # kernel, padding, dilation
113
+ """Pad to 'same' shape outputs."""
114
+ if d > 1:
115
+ k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k] # actual kernel-size
116
+ if p is None:
117
+ p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad
118
+ return p
119
+
120
+
121
+ #######################
122
+ ## Codebase from YOLOv8
123
+ ## END
124
+ #######################
125
+
126
+ def pixel_unshuffle(data, factor=2):
127
+ # performs nn.PixelShuffle(factor) in reverse, torch has some bug for ONNX and TRT, so doing it manually
128
+ B, C, H, W = data.shape
129
+ return data.view(B, C, factor, H//factor, factor, W//factor).permute(0,1,2,4,3,5).reshape(B, -1, H//factor, W//factor)
130
+
131
+ class SwiGLU(nn.Module):
132
+ # should be more advanced, but doesnt improve results so far
133
+ def forward(self, x):
134
+ x, gate = x.chunk(2, dim=-1)
135
+ return F.silu(gate) * x
136
+
137
+
138
+ def window_partition(x, window_size):
139
+ """
140
+ Function for partitioning image into windows and later do windowed attention
141
+ Args:
142
+ x: (B, C, H, W)
143
+ window_size: window size
144
+ Returns:
145
+ windows - local window features (num_windows*B, window_size*window_size, C)
146
+ (Hp, Wp) - the size of the padded image
147
+ """
148
+ B, C, H, W = x.shape
149
+
150
+ if window_size == 0 or (window_size==H and window_size==W):
151
+ windows = x.flatten(2).transpose(1, 2)
152
+ Hp, Wp = H, W
153
+ else:
154
+ pad_h = (window_size - H % window_size) % window_size
155
+ pad_w = (window_size - W % window_size) % window_size
156
+ if pad_h > 0 or pad_w > 0:
157
+ x = F.pad(x, (0, pad_w, 0, pad_h), mode="reflect")
158
+ Hp, Wp = H + pad_h, W + pad_w
159
+
160
+ x = x.view(B, C, Hp // window_size, window_size, Wp // window_size, window_size)
161
+ windows = x.permute(0, 2, 4, 3, 5, 1).reshape(-1, window_size*window_size, C)
162
+
163
+ return windows, (Hp, Wp)
164
+
165
+ class Conv2d_BN(nn.Module):
166
+ '''
167
+ Conv2d + BN layer with folding capability to speed up inference
168
+ Can be merged with Conv() function with additional arguments
169
+ '''
170
+ def __init__(self, a, b, kernel_size=1, stride=1, padding=0, dilation=1, groups=1, bn_weight_init=1, bias=False):
171
+ super().__init__()
172
+ self.conv = torch.nn.Conv2d(a, b, kernel_size, stride, padding, dilation, groups, bias=False)
173
+ if 1:
174
+ self.bn = torch.nn.BatchNorm2d(b)
175
+ torch.nn.init.constant_(self.bn.weight, bn_weight_init)
176
+ torch.nn.init.constant_(self.bn.bias, 0)
177
+
178
+ def forward(self,x):
179
+ x = self.conv(x)
180
+ x = self.bn(x)
181
+ return x
182
+
183
+ @torch.no_grad()
184
+ def switch_to_deploy(self):
185
+ if not isinstance(self.bn, nn.Identity):
186
+ c, bn = self.conv, self.bn
187
+ w = bn.weight / (bn.running_var + bn.eps) ** 0.5
188
+ w = c.weight * w[:, None, None, None]
189
+ b = bn.bias - bn.running_mean * bn.weight / \
190
+ (bn.running_var + bn.eps)**0.5
191
+ self.conv.weight.data.copy_(w)
192
+ self.conv.bias = nn.Parameter(b)
193
+ self.bn = nn.Identity()
194
+
195
+
196
+
197
+ def window_reverse(windows, window_size, H, W, pad_hw):
198
+ """
199
+ Windows to the full feature map
200
+ Args:
201
+ windows: local window features (num_windows*B, window_size, window_size, C)
202
+ window_size: Window size
203
+ H: Height of image
204
+ W: Width of image
205
+ pad_w - a tuple of image passing used in windowing step
206
+ Returns:
207
+ x: (B, C, H, W)
208
+
209
+ """
210
+ # print(f"window_reverse, windows.shape {windows.shape}")
211
+ Hp, Wp = pad_hw
212
+ if window_size == 0 or (window_size==H and window_size==W):
213
+ B = int(windows.shape[0] / (Hp * Wp / window_size / window_size))
214
+ x = windows.transpose(1, 2).view(B, -1, H, W)
215
+ else:
216
+ B = int(windows.shape[0] / (Hp * Wp / window_size / window_size))
217
+ x = windows.view(B, Hp // window_size, Wp // window_size, window_size, window_size, -1)
218
+ x = x.permute(0, 5, 1, 3, 2, 4).reshape(B,windows.shape[2], Hp, Wp)
219
+
220
+ if Hp > H or Wp > W:
221
+ x = x[:, :, :H, :W, ].contiguous()
222
+
223
+ return x
224
+
225
+
226
+
227
+ class PosEmbMLPSwinv2D(nn.Module):
228
+ """
229
+ 2D positional embedding from Swin Transformer v2
230
+ Added functionality to store the positional embedding in the model and not recompute it every time
231
+ """
232
+ def __init__(
233
+ self, window_size, pretrained_window_size, num_heads, seq_length, no_log=False, cpb_mlp_hidden=512,
234
+ ):
235
+ super().__init__()
236
+ self.window_size = window_size
237
+ self.num_heads = num_heads
238
+ # mlp to generate continuous relative position bias
239
+ self.cpb_mlp = nn.Sequential(
240
+ nn.Linear(2, cpb_mlp_hidden, bias=True),
241
+ nn.ReLU(inplace=True),
242
+ nn.Linear(cpb_mlp_hidden, num_heads, bias=False),
243
+ )
244
+
245
+ self.grid_exists = False
246
+ self.seq_length = seq_length
247
+ self.deploy = False
248
+ self.num_heads = num_heads
249
+ self.no_log = no_log
250
+ self.pretrained_window_size = pretrained_window_size
251
+ self.relative_bias_window_size = window_size
252
+
253
+ relative_coords_table, relative_position_index, relative_bias = self.relative_bias_initialization(window_size, num_heads,
254
+ pretrained_window_size, seq_length,
255
+ no_log)
256
+
257
+ self.register_buffer("relative_coords_table", relative_coords_table)
258
+ self.register_buffer("relative_position_index", relative_position_index)
259
+ self.register_buffer("relative_bias", relative_bias) # for EMA
260
+
261
+ def relative_bias_initialization(self, window_size, num_heads, pretrained_window_size, seq_length, no_log):
262
+ # as in separate function to support window size chage after model weights loading
263
+ relative_coords_h = torch.arange(
264
+ -(window_size[0] - 1), window_size[0], dtype=torch.float32
265
+ )
266
+ relative_coords_w = torch.arange(
267
+ -(window_size[1] - 1), window_size[1], dtype=torch.float32
268
+ )
269
+ relative_coords_table = (
270
+ torch.stack(torch.meshgrid([relative_coords_h, relative_coords_w]))
271
+ .permute(1, 2, 0)
272
+ .contiguous()
273
+ .unsqueeze(0)
274
+ ) # 1, 2*Wh-1, 2*Ww-1, 2
275
+ if pretrained_window_size[0] > 0:
276
+ relative_coords_table[:, :, :, 0] /= pretrained_window_size[0] - 1
277
+ relative_coords_table[:, :, :, 1] /= pretrained_window_size[1] - 1
278
+ else:
279
+ relative_coords_table[:, :, :, 0] /= self.window_size[0] - 1
280
+ relative_coords_table[:, :, :, 1] /= self.window_size[1] - 1
281
+
282
+ if not no_log:
283
+ relative_coords_table *= 8 # normalize to -8, 8
284
+ relative_coords_table = (
285
+ torch.sign(relative_coords_table)
286
+ * torch.log2(torch.abs(relative_coords_table) + 1.0)
287
+ / np.log2(8)
288
+ )
289
+
290
+ # get pair-wise relative position index for each token inside the window
291
+ coords_h = torch.arange(self.window_size[0])
292
+ coords_w = torch.arange(self.window_size[1])
293
+ coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww
294
+ coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww
295
+ relative_coords = (
296
+ coords_flatten[:, :, None] - coords_flatten[:, None, :]
297
+ ) # 2, Wh*Ww, Wh*Ww
298
+ relative_coords = relative_coords.permute(
299
+ 1, 2, 0
300
+ ).contiguous() # Wh*Ww, Wh*Ww, 2
301
+ relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0
302
+ relative_coords[:, :, 1] += self.window_size[1] - 1
303
+ relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
304
+ relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww
305
+
306
+ relative_bias = torch.zeros(1, num_heads, seq_length, seq_length)
307
+
308
+ self.relative_bias_window_size = window_size
309
+
310
+ return relative_coords_table, relative_position_index, relative_bias
311
+
312
+
313
+ def switch_to_deploy(self):
314
+ self.deploy = True
315
+ self.grid_exists = True
316
+
317
+ def forward(self, input_tensor):
318
+ # for efficiency, we want this forward to be folded into a single operation (sum)
319
+ # if resolution stays the same, then we dont need to recompute MLP layers
320
+
321
+ if not self.deploy or self.training:
322
+ self.grid_exists = False
323
+
324
+ #compare if all elements in self.window_size list match those in self.relative_bias_window_size
325
+ if not all([self.window_size[i] == self.relative_bias_window_size[i] for i in range(len(self.window_size))]):
326
+ relative_coords_table, relative_position_index, relative_bias = self.relative_bias_initialization(self.window_size, self.num_heads,
327
+ self.pretrained_window_size, self.seq_length,
328
+ self.no_log)
329
+
330
+ self.relative_coords_table = relative_coords_table.to(self.relative_coords_table.device)
331
+ self.relative_position_index = relative_position_index.to(self.relative_position_index.device)
332
+ self.relative_bias = relative_bias.to(self.relative_bias.device)
333
+
334
+ if self.deploy and self.grid_exists:
335
+ input_tensor = input_tensor + self.relative_bias
336
+ return input_tensor
337
+
338
+ if 1:
339
+ self.grid_exists = True
340
+
341
+ relative_position_bias_table = self.cpb_mlp(
342
+ self.relative_coords_table
343
+ ).view(-1, self.num_heads)
344
+ relative_position_bias = relative_position_bias_table[
345
+ self.relative_position_index.view(-1)
346
+ ].view(
347
+ self.window_size[0] * self.window_size[1],
348
+ self.window_size[0] * self.window_size[1],
349
+ -1,
350
+ ) # Wh*Ww,Wh*Ww,nH
351
+
352
+ relative_position_bias = relative_position_bias.permute(
353
+ 2, 0, 1
354
+ ).contiguous() # nH, Wh*Ww, Wh*Ww
355
+ relative_position_bias = 16 * torch.sigmoid(relative_position_bias)
356
+
357
+ self.relative_bias = relative_position_bias.unsqueeze(0)
358
+
359
+ input_tensor = input_tensor + self.relative_bias
360
+ return input_tensor
361
+
362
+
363
+ class GRAAttentionBlock(nn.Module):
364
+ def __init__(self, window_size, dim_in, dim_out,
365
+ num_heads, drop_path=0., qk_scale=None, qkv_bias=False,
366
+ norm_layer=nn.LayerNorm, layer_scale=None,
367
+ use_swiglu=True,
368
+ subsample_ratio=1, dim_ratio=1, conv_base=False,
369
+ do_windowing=True, multi_query=False, use_shift=0,
370
+ cpb_mlp_hidden=512, conv_groups_ratio=0):
371
+ '''
372
+ Global Resolution Attention Block , see README for details
373
+ Attention with subsampling to get a bigger receptive field for attention
374
+ conv_base - use conv2d instead of avgpool2d for downsample / upsample
375
+
376
+
377
+ '''
378
+ super().__init__()
379
+
380
+ self.shift_size=window_size//2 if use_shift else 0
381
+
382
+ self.do_windowing = do_windowing
383
+ self.subsample_ratio = subsample_ratio
384
+
385
+
386
+
387
+ if do_windowing:
388
+ if conv_base:
389
+ self.downsample_op = nn.Conv2d(dim_in, dim_out, kernel_size=subsample_ratio, stride=subsample_ratio) if subsample_ratio > 1 else nn.Identity()
390
+
391
+
392
+ self.downsample_mixer = nn.Identity()
393
+ self.upsample_mixer = nn.Identity()
394
+ self.upsample_op = nn.ConvTranspose2d(dim_in, dim_out, kernel_size=subsample_ratio, stride=subsample_ratio) if subsample_ratio > 1 else nn.Identity()
395
+ else:
396
+ self.downsample_op = nn.AvgPool2d(kernel_size=subsample_ratio, stride=subsample_ratio) if subsample_ratio > 1 else nn.Identity()
397
+ self.downsample_mixer = Conv2d_BN(dim_in, dim_out, kernel_size=1, stride=1) if subsample_ratio > 1 else nn.Identity()
398
+ self.upsample_mixer = nn.Upsample(scale_factor=subsample_ratio, mode='nearest') if subsample_ratio > 1 else nn.Identity()
399
+ self.upsample_op = Conv2d_BN(dim_in, dim_out, kernel_size=1, stride=1, padding=0, bias=False) if subsample_ratio > 1 else nn.Identity()
400
+
401
+
402
+ # in case there is no downsampling conv we want to have it separately
403
+ # will help with information propagation between windows
404
+ if subsample_ratio == 1:
405
+ # conv_groups_ratio=0
406
+ self.pre_conv = Conv2d_BN(dim_in, dim_in, kernel_size=3, stride=1, padding=1, groups=max(1,int(conv_groups_ratio*dim_in)), bias=False)
407
+ # self.pre_conv = nn.Conv2d(dim_in, dim_in, kernel_size=3, stride=1, padding=1, groups=max(1,int(conv_groups_ratio*dim_in)), bias=False)
408
+ # self.pre_conv_act = nn.ReLU6()
409
+ #for simplicity:
410
+ self.pre_conv_act = nn.Identity()
411
+ if conv_groups_ratio == -1:
412
+ self.pre_conv = nn.Identity()
413
+ self.pre_conv_act = nn.Identity()
414
+
415
+ self.window_size = window_size
416
+
417
+ self.norm1 = norm_layer(dim_in)
418
+
419
+ self.attn = WindowAttention(
420
+ dim_in,
421
+ num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
422
+ resolution=window_size,
423
+ seq_length=window_size**2, dim_out=dim_in, multi_query=multi_query,
424
+ shift_size=self.shift_size, cpb_mlp_hidden=cpb_mlp_hidden)
425
+
426
+ self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
427
+
428
+ use_layer_scale = layer_scale is not None and type(layer_scale) in [int, float]
429
+ self.gamma1 = nn.Parameter(layer_scale * torch.ones(dim_in)) if use_layer_scale else 1
430
+
431
+ ### mlp layer
432
+ mlp_ratio = 4
433
+ self.norm2 = norm_layer(dim_in)
434
+ mlp_hidden_dim = int(dim_in * mlp_ratio)
435
+
436
+ activation = nn.GELU if not use_swiglu else SwiGLU
437
+ mlp_hidden_dim = int((4 * dim_in * 1 / 2) / 64) * 64 if use_swiglu else mlp_hidden_dim
438
+
439
+ self.mlp = Mlp(in_features=dim_in, hidden_features=mlp_hidden_dim, act_layer=activation, use_swiglu=use_swiglu)
440
+
441
+ self.gamma2 = nn.Parameter(layer_scale * torch.ones(dim_in)) if layer_scale else 1
442
+ self.drop_path2=DropPath(drop_path) if drop_path > 0. else nn.Identity()
443
+
444
+
445
+ def forward(self, x):
446
+ skip_connection = x
447
+ attn_mask = None
448
+
449
+ # in case there is no downsampling conv we want to have it separately
450
+ # will help with information propagation
451
+ if self.subsample_ratio == 1:
452
+ x = self.pre_conv_act(self.pre_conv(x)) + skip_connection
453
+
454
+ if self.do_windowing:
455
+ # performing windowing if required
456
+ x = self.downsample_op(x)
457
+ x = self.downsample_mixer(x)
458
+
459
+ if self.window_size>0:
460
+ H, W = x.shape[2], x.shape[3]
461
+
462
+ if self.shift_size > 0 and H>self.window_size and W>self.window_size:
463
+ # @swin like cyclic shift, doesnt show better performance
464
+ x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(2, 3))
465
+
466
+ x, pad_hw = window_partition(x, self.window_size)
467
+
468
+ if self.shift_size > 0 and H>self.window_size and W>self.window_size:
469
+ # set atten matrix to have -100 and the top right square
470
+ # attn[:, :, :-self.shift_size, -self.shift_size:] = -100.0
471
+ # calculate attention mask for SW-MSA
472
+ # not used in final version, can be useful for some cases especially for high res
473
+ H, W = pad_hw
474
+ img_mask = torch.zeros((1, H, W, 1), device=x.device) # 1 H W 1
475
+ h_slices = (slice(0, -self.window_size),
476
+ slice(-self.window_size, -self.shift_size),
477
+ slice(-self.shift_size, None))
478
+ w_slices = (slice(0, -self.window_size),
479
+ slice(-self.window_size, -self.shift_size),
480
+ slice(-self.shift_size, None))
481
+ cnt = 0
482
+ for h in h_slices:
483
+ for w in w_slices:
484
+ img_mask[:, h, w, :] = cnt
485
+ cnt += 1
486
+ img_mask = img_mask.transpose(1,2).transpose(1,3)
487
+ mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1
488
+
489
+ mask_windows = mask_windows[0].view(-1, self.window_size * self.window_size)
490
+ attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
491
+ attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
492
+
493
+ # window attention
494
+ x = x + self.drop_path1(self.gamma1*self.attn(self.norm1(x), attn_mask=attn_mask)) # or pass H,W
495
+ # mlp layer
496
+ x = x + self.drop_path2(self.gamma2*self.mlp(self.norm2(x)))
497
+
498
+ if self.do_windowing:
499
+ if self.window_size > 0:
500
+ x = window_reverse(x, self.window_size, H, W, pad_hw)
501
+
502
+ # reverse cyclic shift
503
+ if self.shift_size > 0 and H>self.window_size and W>self.window_size:
504
+ # @swin like cyclic shift, not tested
505
+ x = torch.roll(x, shifts=(self.shift_size, self.shift_size), dims=(2, 3))
506
+
507
+ x = self.upsample_mixer(x)
508
+ x = self.upsample_op(x)
509
+
510
+
511
+ if x.shape[2] != skip_connection.shape[2] or x.shape[3] != skip_connection.shape[3]:
512
+ x = torch.nn.functional.pad(x, ( 0, -x.shape[3] + skip_connection.shape[3], 0, -x.shape[2] + skip_connection.shape[2]), mode="reflect")
513
+ # need to add skip connection because downsampling and upsampling will break residual connection
514
+ # 0.5 is needed to make sure that the skip connection is not too strong
515
+ # in case of no downsample / upsample we can show that 0.5 compensates for the residual connection
516
+ x = 0.5 * x + 0.5 * skip_connection
517
+ return x
518
+
519
+
520
+
521
+
522
+ class MultiResolutionAttention(nn.Module):
523
+ """
524
+ MultiResolutionAttention (MRA) module
525
+ The idea is to use multiple attention blocks with different resolution
526
+ Feature maps are downsampled / upsampled for each attention block on different blocks
527
+ Every attention block supports windowing
528
+ """
529
+
530
+ def __init__(self, window_size, sr_ratio,
531
+ dim, dim_ratio, num_heads,
532
+ do_windowing=True,
533
+ layer_scale=1e-5, norm_layer=nn.LayerNorm,
534
+ drop_path = 0, qkv_bias=False, qk_scale=1.0,
535
+ use_swiglu=True, multi_query=False, conv_base=False,
536
+ use_shift=0, cpb_mlp_hidden=512, conv_groups_ratio=0) -> None:
537
+ """
538
+ Args:
539
+ input_resolution: input image resolution
540
+ window_size: window size
541
+ compression_ratio: compression ratio
542
+ max_depth: maximum depth of the GRA module
543
+ use_shift: do window shifting
544
+ """
545
+ super().__init__()
546
+
547
+ depth = len(sr_ratio)
548
+
549
+ self.attention_blocks = nn.ModuleList()
550
+
551
+
552
+ for i in range(depth):
553
+ subsample_ratio = sr_ratio[i]
554
+ if len(window_size) > i:
555
+ window_size_local = window_size[i]
556
+ else:
557
+ window_size_local = window_size[0]
558
+
559
+ self.attention_blocks.append(GRAAttentionBlock(window_size=window_size_local,
560
+ dim_in=dim, dim_out=dim, num_heads=num_heads,
561
+ qkv_bias=qkv_bias, qk_scale=qk_scale, norm_layer=norm_layer,
562
+ layer_scale=layer_scale, drop_path=drop_path,
563
+ use_swiglu=use_swiglu, subsample_ratio=subsample_ratio, dim_ratio=dim_ratio,
564
+ do_windowing=do_windowing, multi_query=multi_query, conv_base=conv_base,
565
+ use_shift=use_shift, cpb_mlp_hidden=cpb_mlp_hidden, conv_groups_ratio=conv_groups_ratio),
566
+ )
567
+
568
+ def forward(self, x):
569
+
570
+ for attention_block in self.attention_blocks:
571
+ x = attention_block(x)
572
+
573
+ return x
574
+
575
+
576
+
577
+ class Mlp(nn.Module):
578
+ """
579
+ Multi-Layer Perceptron (MLP) block
580
+ """
581
+
582
+ def __init__(self,
583
+ in_features,
584
+ hidden_features=None,
585
+ out_features=None,
586
+ act_layer=nn.GELU,
587
+ use_swiglu=True,
588
+ drop=0.):
589
+ """
590
+ Args:
591
+ in_features: input features dimension.
592
+ hidden_features: hidden features dimension.
593
+ out_features: output features dimension.
594
+ act_layer: activation function.
595
+ drop: dropout rate.
596
+ """
597
+
598
+ super().__init__()
599
+ out_features = out_features or in_features
600
+ hidden_features = hidden_features or in_features
601
+ self.fc1 = nn.Linear(in_features, hidden_features * (2 if use_swiglu else 1), bias=False)
602
+ self.act = act_layer()
603
+ self.fc2 = nn.Linear(hidden_features, out_features, bias=False)
604
+
605
+ def forward(self, x):
606
+ x_size = x.size()
607
+ x = x.view(-1, x_size[-1])
608
+ x = self.fc1(x)
609
+ x = self.act(x)
610
+ x = self.fc2(x)
611
+ x = x.view(x_size)
612
+ return x
613
+
614
+ class Downsample(nn.Module):
615
+ """
616
+ Down-sampling block
617
+ Pixel Unshuffle is used for down-sampling, works great accuracy - wise but takes 10% more TRT time
618
+ """
619
+
620
+ def __init__(self,
621
+ dim,
622
+ shuffle = False,
623
+ ):
624
+ """
625
+ Args:
626
+ dim: feature size dimension.
627
+ shuffle: idea with
628
+ keep_dim: bool argument for maintaining the resolution.
629
+ """
630
+
631
+ super().__init__()
632
+ dim_out = 2 * dim
633
+
634
+ if shuffle:
635
+ self.norm = lambda x: pixel_unshuffle(x, factor=2)
636
+ self.reduction = Conv2d_BN(dim*4, dim_out, 1, 1, 0, bias=False)
637
+ # pixel unshuffleging works well but doesnt provide any speedup
638
+ else:
639
+ # removed layer norm for better, in this formulation we are getting 10% better speed
640
+ # LayerNorm for high resolution inputs will be a pain as it pools over the entire spatial dimension
641
+ # therefore we remove it compared to the original implementation in FasterViT
642
+ self.norm = nn.Identity()
643
+ self.reduction = Conv2d_BN(dim, dim_out, 3, 2, 1, bias=False)
644
+
645
+
646
+ def forward(self, x):
647
+ x = self.norm(x)
648
+ x = self.reduction(x)
649
+ return x
650
+
651
+
652
+ class PatchEmbed(nn.Module):
653
+ """
654
+ Patch embedding block
655
+ Used to convert image into an initial set of feature maps with lower resolution
656
+ """
657
+
658
+ def __init__(self, in_chans=3, in_dim=64, dim=96, shuffle_down=False):
659
+ """
660
+ Args:
661
+ in_chans: number of input channels.
662
+ in_dim: intermediate feature size dimension to speed up stem.
663
+ dim: final stem channel number
664
+ shuffle_down: use PixelUnshuffle for down-sampling, effectively increases the receptive field
665
+ """
666
+
667
+ super().__init__()
668
+ # shuffle_down = False
669
+ if not shuffle_down:
670
+ self.proj = nn.Identity()
671
+ self.conv_down = nn.Sequential(
672
+ Conv2d_BN(in_chans, in_dim, 3, 2, 1, bias=False),
673
+ nn.ReLU(),
674
+ Conv2d_BN(in_dim, dim, 3, 2, 1, bias=False),
675
+ nn.ReLU()
676
+ )
677
+ else:
678
+ self.proj = lambda x: pixel_unshuffle(x, factor=4)
679
+ self.conv_down = nn.Sequential(Conv2d_BN(in_chans*16, dim, 3, 1, 1),
680
+ nn.ReLU(),
681
+ )
682
+
683
+ def forward(self, x):
684
+ x = self.proj(x)
685
+ x = self.conv_down(x)
686
+ return x
687
+
688
+
689
+
690
+ class ConvBlock(nn.Module):
691
+ """
692
+ Convolutional block, used in first couple of stages
693
+ Experimented with plan resnet-18 like modules, they are the best in terms of throughput
694
+ Finally, YOLOv8 idea seem to work fine (resnet-18 like block with squeezed feature dimension, and feature concatendation at the end)
695
+ """
696
+ def __init__(self, dim,
697
+ drop_path=0.,
698
+ layer_scale=None,
699
+ kernel_size=3,
700
+ ):
701
+ super().__init__()
702
+
703
+ self.conv1 = Conv2d_BN(dim, dim, kernel_size=kernel_size, stride=1, padding=1)
704
+ self.act1 = nn.GELU()
705
+
706
+ self.conv2 = Conv2d_BN(dim, dim, kernel_size=kernel_size, stride=1, padding=1)
707
+
708
+ self.layer_scale = layer_scale
709
+ if layer_scale is not None and type(layer_scale) in [int, float]:
710
+ self.gamma = nn.Parameter(layer_scale * torch.ones(dim))
711
+ self.layer_scale = True
712
+ else:
713
+ self.layer_scale = False
714
+ self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
715
+
716
+ def forward(self, x):
717
+ input = x
718
+
719
+ x = self.conv1(x)
720
+ x = self.act1(x)
721
+ x = self.conv2(x)
722
+
723
+ if self.layer_scale:
724
+ x = x * self.gamma.view(1, -1, 1, 1)
725
+ x = input + self.drop_path(x)
726
+ return x
727
+
728
+
729
+ class WindowAttention(nn.Module):
730
+ # Windowed Attention from SwinV2
731
+ # use a MLP trick to deal with various input image resolutions, then fold it to improve speed
732
+
733
+ def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, resolution=0,
734
+ seq_length=0, dim_out=None, multi_query=False, shift_size=0, cpb_mlp_hidden=512):
735
+ # taken from EdgeViT and tweaked with attention bias.
736
+ super().__init__()
737
+ if not dim_out: dim_out = dim
738
+ self.shift_size = shift_size
739
+ self.multi_query = multi_query
740
+ self.num_heads = num_heads
741
+ head_dim = dim // num_heads
742
+ self.head_dim = dim // num_heads
743
+
744
+ self.dim_internal = dim
745
+
746
+ self.scale = qk_scale or head_dim ** -0.5
747
+ if not multi_query:
748
+ self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
749
+ else:
750
+ self.qkv = nn.Linear(dim, dim + 2*self.head_dim, bias=qkv_bias)
751
+
752
+ self.proj = nn.Linear(dim, dim_out, bias=False)
753
+ # attention positional bias
754
+ self.pos_emb_funct = PosEmbMLPSwinv2D(window_size=[resolution, resolution],
755
+ pretrained_window_size=[resolution, resolution],
756
+ num_heads=num_heads,
757
+ seq_length=seq_length,
758
+ cpb_mlp_hidden=cpb_mlp_hidden)
759
+
760
+ self.resolution = resolution
761
+
762
+ def forward(self, x, attn_mask = None):
763
+ B, N, C = x.shape
764
+
765
+ if not self.multi_query:
766
+ qkv = self.qkv(x).reshape(B, -1, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
767
+ q, k, v = qkv[0], qkv[1], qkv[2]
768
+ else:
769
+ qkv = self.qkv(x)
770
+ (q, k, v) = qkv.split([self.dim_internal, self.head_dim, self.head_dim], dim=2)
771
+
772
+ q = q.reshape(B, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
773
+ k = k.reshape(B, -1, 1, C // self.num_heads).permute(0, 2, 1, 3)
774
+ v = v.reshape(B, -1, 1, C // self.num_heads).permute(0, 2, 1, 3)
775
+
776
+ attn = (q @ k.transpose(-2, -1)) * self.scale
777
+
778
+ attn = self.pos_emb_funct(attn)
779
+
780
+ #add window shift
781
+ if attn_mask is not None:
782
+ nW = attn_mask.shape[0]
783
+ attn = attn.view(B // nW, nW, self.num_heads, N, N) + attn_mask.unsqueeze(1).unsqueeze(0)
784
+ attn = attn.view(-1, self.num_heads, N, N)
785
+
786
+ attn = attn.softmax(dim=-1)
787
+ x = (attn @ v).transpose(1, 2).reshape(B, -1, C)
788
+ x = self.proj(x)
789
+ return x
790
+
791
+
792
+
793
+ class ERADIOLayer(nn.Module):
794
+ """
795
+ E-RADIO Layer
796
+ """
797
+
798
+ def __init__(self,
799
+ dim,
800
+ depth,
801
+ num_heads,
802
+ window_size,
803
+ conv=False,
804
+ downsample=True,
805
+ mlp_ratio=4.,
806
+ qkv_bias=False,
807
+ qk_scale=None,
808
+ norm_layer=nn.LayerNorm,
809
+ drop_path=0.,
810
+ layer_scale=None,
811
+ layer_scale_conv=None,
812
+ sr_dim_ratio=1,
813
+ sr_ratio=1,
814
+ multi_query=False,
815
+ use_swiglu=True,
816
+ yolo_arch=False,
817
+ downsample_shuffle=False,
818
+ conv_base=False,
819
+ use_shift=False,
820
+ cpb_mlp_hidden=512,
821
+ conv_groups_ratio=0,
822
+ verbose: bool = True,
823
+
824
+ ):
825
+ """
826
+ Args:
827
+ dim: feature size dimension.
828
+ depth: number of layers in each stage.
829
+ input_resolution: input image resolution.
830
+ window_size: window size in each stage.
831
+ downsample: bool argument for down-sampling.
832
+ mlp_ratio: MLP ratio.
833
+ num_heads: number of heads in each stage.
834
+ qkv_bias: bool argument for query, key, value learnable bias.
835
+ qk_scale: bool argument to scaling query, key.
836
+ drop: dropout rate.
837
+ attn_drop: attention dropout rate.
838
+ drop_path: drop path rate.
839
+ norm_layer: normalization layer.
840
+ layer_scale: layer scaling coefficient.
841
+ use_shift: SWIN like window shifting for half the window size for every alternating layer (considering multi-resolution)
842
+ conv_groups_ratio: group ratio for conv when no subsampling in multi-res attention
843
+ """
844
+
845
+ super().__init__()
846
+ self.conv = conv
847
+ self.yolo_arch=False
848
+ self.verbose = verbose
849
+ if conv:
850
+ if not yolo_arch:
851
+ self.blocks = nn.ModuleList([
852
+ ConvBlock(dim=dim,
853
+ drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
854
+ layer_scale=layer_scale_conv)
855
+ for i in range(depth)])
856
+ self.blocks = nn.Sequential(*self.blocks)
857
+ else:
858
+ self.blocks = C2f(dim,dim,n=depth,shortcut=True,e=0.5)
859
+ self.yolo_arch=True
860
+ else:
861
+ if not isinstance(window_size, list): window_size = [window_size]
862
+ self.window_size = window_size[0]
863
+ self.do_single_windowing = True
864
+ if not isinstance(sr_ratio, list): sr_ratio = [sr_ratio]
865
+ self.sr_ratio = sr_ratio
866
+ if any([sr!=1 for sr in sr_ratio]) or len(set(window_size))>1:
867
+ self.do_single_windowing = False
868
+ do_windowing = True
869
+ else:
870
+ self.do_single_windowing = True
871
+ do_windowing = False
872
+
873
+ #for v2_2
874
+ if conv_groups_ratio != -1:
875
+ self.do_single_windowing = False
876
+ do_windowing = True
877
+
878
+ self.blocks = nn.ModuleList()
879
+ for i in range(depth):
880
+ self.blocks.append(
881
+ MultiResolutionAttention(window_size=window_size,
882
+ sr_ratio=sr_ratio,
883
+ dim=dim,
884
+ dim_ratio = sr_dim_ratio,
885
+ num_heads=num_heads,
886
+ norm_layer=norm_layer,
887
+ drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
888
+ layer_scale=layer_scale,
889
+ qkv_bias=qkv_bias,
890
+ qk_scale=qk_scale,
891
+ use_swiglu=use_swiglu,
892
+ do_windowing=do_windowing,
893
+ multi_query=multi_query,
894
+ conv_base=conv_base,
895
+ cpb_mlp_hidden=cpb_mlp_hidden,
896
+ use_shift =0 if ((not use_shift) or ((i) % 2 == 0)) else True ,
897
+ conv_groups_ratio=conv_groups_ratio,
898
+ ))
899
+ self.blocks = nn.Sequential(*self.blocks)
900
+
901
+ self.transformer = not conv
902
+ self.downsample = None if not downsample else Downsample(dim=dim, shuffle=downsample_shuffle)
903
+
904
+
905
+ def forward(self, x):
906
+ B, C, H, W = x.shape
907
+
908
+ # do padding for transforemr
909
+ interpolate = True
910
+ if self.transformer and interpolate:
911
+ # Windowed Attention will split feature map into windows with the size of window_size x window_size
912
+ # if the resolution is not divisible by window_size, we need to interpolate the feature map
913
+ # can be done via padding, but doing so after training hurts the model performance.
914
+ # interpolation affects the performance as well, but not as much as padding
915
+ if isinstance(self.window_size, list) or isinstance(self.window_size, tuple):
916
+ current_max_window_size = max(self.window_size)
917
+ else:
918
+ current_max_window_size = self.window_size
919
+
920
+ max_window_size = max([res_upsample*current_max_window_size for res_upsample in self.sr_ratio])
921
+ if H % max_window_size != 0 or W % max_window_size != 0:
922
+ new_h = int(np.ceil(H/max_window_size)*max_window_size)
923
+ new_w = int(np.ceil(W/max_window_size)*max_window_size)
924
+ x = F.interpolate(x, size=(new_h, new_w), mode='nearest')
925
+ if self.verbose:
926
+ warnings.warn(f"Choosen window size is not optimal for given resolution. Interpolation of features maps will be done and it can affect the performance. Max window size is {max_window_size}, feature map size is {H}x{W}, interpolated feature map size is {new_h}x{new_w}.")
927
+
928
+
929
+ if self.transformer and self.do_single_windowing:
930
+ H, W = x.shape[2], x.shape[3]
931
+ x, pad_hw = window_partition(x, self.window_size)
932
+
933
+ #run main blocks
934
+ x = self.blocks(x)
935
+
936
+ if self.transformer and self.do_single_windowing:
937
+ x = window_reverse(x, self.window_size, H, W, pad_hw)
938
+
939
+ if self.transformer and interpolate:
940
+ #lets keep original resolution, might be not ideal, but for the upsampling tower we need to keep the expected resolution.
941
+ x = F.interpolate(x, size=(H, W), mode='nearest')
942
+
943
+ if self.downsample is None:
944
+ return x, x
945
+
946
+ return self.downsample(x), x # changing to output pre downsampled features
947
+
948
+
949
+ class InterpolateLayer(nn.Module):
950
+ def __init__(self, size=None, scale_factor=None, mode='nearest'):
951
+ super(InterpolateLayer, self).__init__()
952
+ self.size = size
953
+ self.scale_factor = scale_factor
954
+ self.mode = mode
955
+
956
+ def forward(self, x):
957
+ return F.interpolate(x, size=self.size, scale_factor=self.scale_factor, mode=self.mode)
958
+
959
+
960
+ class HiResNeck(nn.Module):
961
+ """
962
+ The block is used to output dense features from all stages
963
+ Otherwise, by default, only the last stage features are returned with E-RADIO
964
+ """
965
+ def __init__(self, dim, depths, neck_start_stage, full_features_head_dim, downsample_enabled):
966
+
967
+ '''
968
+ Hi Resolution neck to support output of high res features that are useful for dense tasks.
969
+ depths - total number of layers in the base model
970
+ neck_start_stage - when to start the neck, 0 - start from the first stage, 1 - start from the second stage etc.
971
+ earlier layers result in higher resolution features at the cost of compute
972
+ full_features_head_dim - number of channels in the dense features head
973
+ '''
974
+ super().__init__()
975
+ # create feature projection layers for segmentation output
976
+ self.neck_features_proj = nn.ModuleList()
977
+ self.neck_start_stage = neck_start_stage
978
+ upsample_ratio = 1
979
+ for i in range(len(depths)):
980
+ level_n_features_output = int(dim * 2 ** i)
981
+
982
+ if self.neck_start_stage > i: continue
983
+
984
+ if (upsample_ratio > 1) or full_features_head_dim!=level_n_features_output:
985
+ feature_projection = nn.Sequential()
986
+ if False:
987
+ feature_projection.add_module("norm",nn.BatchNorm2d(level_n_features_output)) #fast, but worse
988
+ feature_projection.add_module("dconv", nn.ConvTranspose2d(level_n_features_output,
989
+ full_features_head_dim, kernel_size=upsample_ratio, stride=upsample_ratio))
990
+ else:
991
+ # B, in_channels, H, W -> B, in_channels, H*upsample_ratio, W*upsample_ratio
992
+ # print("upsample ratio", upsample_ratio, level_n_features_output, level_n_features_output)
993
+ feature_projection.add_module("upsample", InterpolateLayer(scale_factor=upsample_ratio, mode='nearest'))
994
+ feature_projection.add_module("conv1", nn.Conv2d(level_n_features_output, level_n_features_output, kernel_size=3, stride=1, padding=1, groups=level_n_features_output))
995
+ feature_projection.add_module("norm",nn.BatchNorm2d(level_n_features_output))
996
+ # B, in_channels, H*upsample_ratio, W*upsample_ratio -> B, full_features_head_dim, H*upsample_ratio, W*upsample_ratio
997
+ feature_projection.add_module("conv2", nn.Conv2d(level_n_features_output, full_features_head_dim, kernel_size=1, stride=1, padding=0))
998
+ else:
999
+ feature_projection = nn.Sequential()
1000
+
1001
+ self.neck_features_proj.append(feature_projection)
1002
+
1003
+ if i>0 and downsample_enabled[i]:
1004
+ upsample_ratio *= 2
1005
+
1006
+ def forward(self, x, il_level=-1, full_features=None):
1007
+ if self.neck_start_stage > il_level:
1008
+ return full_features
1009
+
1010
+ if full_features is None:
1011
+ full_features = self.neck_features_proj[il_level - self.neck_start_stage](x)
1012
+ else:
1013
+ #upsample torch tensor x to match full_features size, and add to full_features
1014
+ feature_projection = self.neck_features_proj[il_level - self.neck_start_stage](x)
1015
+ if feature_projection.shape[2] != full_features.shape[2] or feature_projection.shape[3] != full_features.shape[3]:
1016
+ feature_projection = torch.nn.functional.pad(feature_projection, ( 0, -feature_projection.shape[3] + full_features.shape[3], 0, -feature_projection.shape[2] + full_features.shape[2]))
1017
+ full_features = full_features + feature_projection
1018
+ return full_features
1019
+
1020
+ class ERADIO(nn.Module):
1021
+ """
1022
+ Efficient RADIO
1023
+ """
1024
+
1025
+ def __init__(self,
1026
+ dim,
1027
+ in_dim,
1028
+ depths,
1029
+ window_size,
1030
+ mlp_ratio,
1031
+ num_heads,
1032
+ drop_path_rate=0.2,
1033
+ in_chans=3,
1034
+ num_classes=1000,
1035
+ qkv_bias=False,
1036
+ qk_scale=None,
1037
+ layer_scale=None,
1038
+ layer_scale_conv=None,
1039
+ layer_norm_last=False,
1040
+ sr_ratio = [1, 1, 1, 1],
1041
+ max_depth = -1,
1042
+ conv_base=False,
1043
+ use_swiglu=False,
1044
+ multi_query=False,
1045
+ norm_layer=nn.LayerNorm,
1046
+ drop_uniform=False,
1047
+ yolo_arch=False,
1048
+ shuffle_down=False,
1049
+ downsample_shuffle=False,
1050
+ return_full_features=False,
1051
+ full_features_head_dim=128,
1052
+ neck_start_stage=1,
1053
+ use_neck=False,
1054
+ use_shift=False,
1055
+ cpb_mlp_hidden=512,
1056
+ conv_groups_ratio=0,
1057
+ verbose: bool = False,
1058
+ **kwargs):
1059
+ """
1060
+ Args:
1061
+ dim: feature size dimension.
1062
+ depths: number of layers in each stage.
1063
+ window_size: window size in each stage.
1064
+ mlp_ratio: MLP ratio.
1065
+ num_heads: number of heads in each stage.
1066
+ drop_path_rate: drop path rate.
1067
+ in_chans: number of input channels.
1068
+ num_classes: number of classes.
1069
+ qkv_bias: bool argument for query, key, value learnable bias.
1070
+ qk_scale: bool argument to scaling query, key.
1071
+ drop_rate: dropout rate.
1072
+ attn_drop_rate: attention dropout rate.
1073
+ norm_layer: normalization layer.
1074
+ layer_scale: layer scaling coefficient.
1075
+ return_full_features: output dense features as well as logits
1076
+ full_features_head_dim: number of channels in the dense features head
1077
+ neck_start_stage: a stage id to start full feature neck. Model has 4 stages, indix starts with 0
1078
+ for 224 resolution, the output of the stage before downsample:
1079
+ stage 0: 56x56, stage 1: 28x28, stage 2: 14x14, stage 3: 7x7
1080
+ use_neck: even for summarization embedding use neck
1081
+ use_shift: SWIN like window shifting but without masking attention
1082
+ conv_groups_ratio: will be used for conv blocks where there is no multires attention,
1083
+ if 0 then normal conv,
1084
+ if 1 then channels are independent,
1085
+ if -1 then no conv at all
1086
+
1087
+ """
1088
+ super().__init__()
1089
+
1090
+ num_features = int(dim * 2 ** (len(depths) - 1))
1091
+ self.num_classes = num_classes
1092
+ self.patch_embed = PatchEmbed(in_chans=in_chans, in_dim=in_dim, dim=dim, shuffle_down=shuffle_down)
1093
+ # set return_full_features true if we want to return full features from all stages
1094
+ self.return_full_features = return_full_features
1095
+ self.use_neck = use_neck
1096
+
1097
+ dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
1098
+ if drop_uniform:
1099
+ dpr = [drop_path_rate for x in range(sum(depths))]
1100
+
1101
+ if not isinstance(max_depth, list): max_depth = [max_depth] * len(depths)
1102
+
1103
+ self.levels = nn.ModuleList()
1104
+ for i in range(len(depths)):
1105
+ conv = True if (i == 0 or i == 1) else False
1106
+
1107
+ level = ERADIOLayer(dim=int(dim * 2 ** i),
1108
+ depth=depths[i],
1109
+ num_heads=num_heads[i],
1110
+ window_size=window_size[i],
1111
+ mlp_ratio=mlp_ratio,
1112
+ qkv_bias=qkv_bias,
1113
+ qk_scale=qk_scale,
1114
+ conv=conv,
1115
+ drop_path=dpr[sum(depths[:i]):sum(depths[:i + 1])],
1116
+ downsample=(i < len(depths) - 1),
1117
+ layer_scale=layer_scale,
1118
+ layer_scale_conv=layer_scale_conv,
1119
+ sr_ratio=sr_ratio[i],
1120
+ use_swiglu=use_swiglu,
1121
+ multi_query=multi_query,
1122
+ norm_layer=norm_layer,
1123
+ yolo_arch=yolo_arch,
1124
+ downsample_shuffle=downsample_shuffle,
1125
+ conv_base=conv_base,
1126
+ cpb_mlp_hidden=cpb_mlp_hidden,
1127
+ use_shift=use_shift,
1128
+ conv_groups_ratio=conv_groups_ratio,
1129
+ verbose=verbose)
1130
+
1131
+ self.levels.append(level)
1132
+
1133
+ if self.return_full_features or self.use_neck:
1134
+ #num_heads
1135
+ downsample_enabled = [self.levels[i-1].downsample is not None for i in range(len(self.levels))]
1136
+ self.high_res_neck = HiResNeck(dim, depths, neck_start_stage, full_features_head_dim, downsample_enabled)
1137
+
1138
+ self.switched_to_deploy = False
1139
+
1140
+ self.norm = LayerNorm2d(num_features) if layer_norm_last else nn.BatchNorm2d(num_features)
1141
+ self.avgpool = nn.AdaptiveAvgPool2d(1)
1142
+ self.head = nn.Linear(num_features, num_classes) if num_classes > 0 else nn.Identity()
1143
+ self.apply(self._init_weights)
1144
+
1145
+ def _init_weights(self, m):
1146
+ if isinstance(m, nn.Linear):
1147
+ trunc_normal_(m.weight, std=.02)
1148
+ if isinstance(m, nn.Linear) and m.bias is not None:
1149
+ nn.init.constant_(m.bias, 0)
1150
+ elif isinstance(m, nn.LayerNorm):
1151
+ nn.init.constant_(m.bias, 0)
1152
+ nn.init.constant_(m.weight, 1.0)
1153
+ elif isinstance(m, LayerNorm2d):
1154
+ nn.init.constant_(m.bias, 0)
1155
+ nn.init.constant_(m.weight, 1.0)
1156
+ elif isinstance(m, nn.BatchNorm2d):
1157
+ nn.init.ones_(m.weight)
1158
+ nn.init.zeros_(m.bias)
1159
+
1160
+ @torch.jit.ignore
1161
+ def no_weight_decay_keywords(self):
1162
+ return {'rpb'}
1163
+
1164
+ def forward_features(self, x):
1165
+ _, _, H, W = x.shape
1166
+ if H % 32 != 0 or W % 32 != 0:
1167
+ raise ValueError(f"E-RADIO requires input dimensions to be divisible by 32 but got H x W: {H} x {W}")
1168
+ x = self.patch_embed(x)
1169
+ full_features = None
1170
+ for il, level in enumerate(self.levels):
1171
+ x, pre_downsample_x = level(x)
1172
+
1173
+ if self.return_full_features or self.use_neck:
1174
+ full_features = self.high_res_neck(pre_downsample_x, il, full_features)
1175
+
1176
+ # x = self.norm(full_features if (self.return_full_features or self.use_neck) else x)
1177
+ x = self.norm(x) # new version for
1178
+
1179
+ if not self.return_full_features:
1180
+ return x, None
1181
+
1182
+ return x, full_features
1183
+
1184
+ def forward(self, x):
1185
+ x, full_features = self.forward_features(x)
1186
+
1187
+ x = self.avgpool(x)
1188
+ x = torch.flatten(x, 1)
1189
+
1190
+ x = self.head(x)
1191
+ if full_features is not None:
1192
+ return x, full_features
1193
+ return x
1194
+
1195
+ def switch_to_deploy(self):
1196
+ '''
1197
+ A method to perform model self-compression
1198
+ merges BN into conv layers
1199
+ converts MLP relative positional bias into precomputed buffers
1200
+ '''
1201
+ if not self.switched_to_deploy:
1202
+ for level in [self.patch_embed, self.levels, self.head]:
1203
+ for module in level.modules():
1204
+ if hasattr(module, 'switch_to_deploy'):
1205
+ module.switch_to_deploy()
1206
+ self.switched_to_deploy = True
1207
+
1208
+
1209
+ def change_window_size(self, new_window_size):
1210
+ """
1211
+ E-RADIO employs windowed attention, which may be sensitive to the choice of this parameter,
1212
+ especially in cases of uneven partitioning of the feature maps.
1213
+ E-RADIO allows for the adjustment of the window size after training,
1214
+ making it adaptable to different input image resolutions.
1215
+ The recommended values for window size based on input resolution are as follows:
1216
+
1217
+ Input Resolution | Window Size
1218
+ 224 | 7
1219
+ 256 | 8
1220
+ 386 | 12
1221
+ 512 | 16
1222
+ Ideally, the window size should be a factor of the input resolution. In the third stage, we divide the resolution by 16, so the window size should be
1223
+ img_res/16/2
1224
+ for the third stage and img_res/32 for the last stage. While this can be applied in a brute-force manner, a better way is to do model.change_window_size.
1225
+ Manual way to change resolution -> model.change_window_size(resolution)
1226
+ """
1227
+ window_size = new_window_size
1228
+ print(f"Setting window size to {window_size}")
1229
+ for module in self.modules():
1230
+ if hasattr(module, "window_size"):
1231
+ # check if tuple or a number
1232
+ if isinstance(module.window_size, tuple):
1233
+ if module.window_size[0] != window_size:
1234
+ module.window_size = (window_size, window_size)
1235
+ elif isinstance(module.window_size, list):
1236
+ if module.window_size[0] != window_size:
1237
+ module.window_size = [window_size, window_size]
1238
+ else:
1239
+ module.window_size = window_size
1240
+
1241
+
1242
+ def set_optimal_window_size(self, image_dim, max_window_size = 16):
1243
+ """
1244
+ Using hand picked window size for various resolutions.
1245
+
1246
+ E-RADIO employs windowed attention, which may be sensitive to the choice of this parameter,
1247
+ especially in cases of uneven partitioning of the feature maps.
1248
+ E-RADIO allows for the adjustment of the window size after training,
1249
+ making it adaptable to different input image resolutions.
1250
+ The recommended values for window size based on input resolution are as follows:
1251
+
1252
+ Input Resolution | Window Size
1253
+ 224 | 7
1254
+ 256 | 8
1255
+ 386 | 12
1256
+ 512 | 16
1257
+ Ideally, the window size should be a factor of the input resolution. In the third stage, we divide the resolution by 16, so the window size should be
1258
+ img_res/16/2
1259
+ for the third stage and img_res/32 for the last stage. While this can be applied in a brute-force manner, a better way is to do model.change_window_size.
1260
+ Manual way to change resolution -> model.change_window_size(resolution)
1261
+
1262
+ """
1263
+ # import math
1264
+
1265
+ def divisorGenerator(n):
1266
+ large_divisors = []
1267
+ for i in range(1, int(math.sqrt(n) + 1)):
1268
+ if n % i == 0:
1269
+ yield i
1270
+ if i*i != n:
1271
+ large_divisors.append(n / i)
1272
+ for divisor in reversed(large_divisors):
1273
+ yield divisor
1274
+
1275
+ if isinstance(image_dim, list) or isinstance(image_dim, tuple):
1276
+ image_dim = min(image_dim)
1277
+
1278
+ # we do windowed attention in the 3rd stage for the first time, therefore //16,
1279
+ # we do subsampled attention with downsample by 2 so need to get //32 actually
1280
+ # ideally we should rewrite this to be dependent on the structure of the model like what if subsampled is removed etc
1281
+ all_divisors = np.array(list(divisorGenerator(image_dim//32)))
1282
+ new_window_size = int(min(all_divisors[all_divisors <= max_window_size][-1], max_window_size))
1283
+
1284
+ # for image_dim in [128, 224, 256, 384, 512, 768, 1024]:
1285
+ # all_divisors = np.array(list(divisorGenerator(image_dim//32)))
1286
+ # new_window_size = int(min(all_divisors[all_divisors <= max_window_size][-1], max_window_size))
1287
+ # print(f"Setting window size to {new_window_size} for image resolution {image_dim}")
1288
+
1289
+ self.change_window_size(new_window_size = new_window_size)
1290
+
1291
+
1292
+ @register_model
1293
+ def eradio_large_fullres_ws16(pretrained=False, **kwargs):
1294
+ model = ERADIO(
1295
+ depths=[3, 3, 5, 5],
1296
+ num_heads=[2, 4, 8, 16],
1297
+ window_size=[None, None, [16, 16], 16],
1298
+ dim=192,
1299
+ in_dim=64,
1300
+ mlp_ratio=4,
1301
+ drop_path_rate=0.0,
1302
+ sr_ratio=[1, 1, [2, 1], 1],
1303
+ use_swiglu=False,
1304
+ yolo_arch=True,
1305
+ shuffle_down=False,
1306
+ conv_base=True,
1307
+ use_neck=True,
1308
+ full_features_head_dim=1536,
1309
+ neck_start_stage=2,
1310
+ **kwargs,
1311
+ )
1312
+ if pretrained:
1313
+ model.load_state_dict(torch.load(pretrained)["state_dict"])
1314
+ return model
1315
+
1316
+
1317
+ @register_model
1318
+ def eradio_xxxtiny(pretrained=False, **kwargs): # ,
1319
+ model = ERADIO(
1320
+ depths=[1, 3, 4, 5],
1321
+ num_heads=[2, 4, 8, 16],
1322
+ window_size=[None, None, [16, 16], 16],
1323
+ dim=32,
1324
+ in_dim=32,
1325
+ mlp_ratio=4,
1326
+ drop_path_rate=0.0,
1327
+ sr_ratio=[1, 1, [2, 1], 1],
1328
+ use_swiglu=False,
1329
+ yolo_arch=True,
1330
+ shuffle_down=False,
1331
+ conv_base=True,
1332
+ use_neck=True,
1333
+ full_features_head_dim=256,
1334
+ neck_start_stage=2,
1335
+ **kwargs,
1336
+ )
1337
+ if pretrained:
1338
+ model.load_state_dict(torch.load(pretrained))
1339
+ return model
1340
+
1341
+ @register_model
1342
+ def eradio_xxxtiny_8x_ws12(pretrained=False, **kwargs):
1343
+ model = ERADIO(depths=[1, 3, 4, 5],
1344
+ num_heads=[2, 4, 8, 16],
1345
+ window_size=[None, None, [12, 12], 12],
1346
+ dim=32,
1347
+ in_dim=32,
1348
+ mlp_ratio=4,
1349
+ drop_path_rate=0.0,
1350
+ sr_ratio=[1, 1, [2, 1], 1],
1351
+ use_swiglu=False,
1352
+ downsample_shuffle=False,
1353
+ yolo_arch=True,
1354
+ shuffle_down=False,
1355
+ cpb_mlp_hidden=64,
1356
+ use_neck=True,
1357
+ full_features_head_dim=256,
1358
+ neck_start_stage=2,
1359
+ conv_groups_ratio = 1,
1360
+ **kwargs)
1361
+ if pretrained:
1362
+ model.load_state_dict(torch.load(pretrained)["state_dict"])
1363
+ return model
1364
+
1365
+
1366
+ @register_model
1367
+ def eradio_xxxtiny_8x_ws16(pretrained=False, **kwargs):
1368
+ model = ERADIO(depths=[1, 3, 4, 5],
1369
+ num_heads=[2, 4, 8, 16],
1370
+ window_size=[None, None, [16, 16], 16],
1371
+ dim=32,
1372
+ in_dim=32,
1373
+ mlp_ratio=4,
1374
+ drop_path_rate=0.0,
1375
+ sr_ratio=[1, 1, [2, 1], 1],
1376
+ use_swiglu=False,
1377
+ downsample_shuffle=False,
1378
+ yolo_arch=True,
1379
+ shuffle_down=False,
1380
+ cpb_mlp_hidden=64,
1381
+ use_neck=True,
1382
+ full_features_head_dim=256,
1383
+ neck_start_stage=1,
1384
+ conv_groups_ratio = 1,
1385
+ **kwargs)
1386
+ if pretrained:
1387
+ model.load_state_dict(torch.load(pretrained)["state_dict"])
1388
+ return model
1389
+
1390
+ @register_model
1391
+ def eradio(pretrained=False, **kwargs):
1392
+ return eradio_large_fullres_ws16(pretrained=pretrained, **kwargs)
extra_timm_models.py ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved.
2
+ #
3
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
4
+ # and proprietary rights in and to this software, related documentation
5
+ # and any modifications thereto. Any use, reproduction, disclosure or
6
+ # distribution of this software and related documentation without an express
7
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
8
+
9
+ from torch import nn
10
+
11
+ from timm.models import register_model
12
+ from timm.models.vision_transformer import VisionTransformer, _create_vision_transformer, Mlp
13
+
14
+
15
+ @register_model
16
+ def vit_tiny_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:
17
+ """ ViT-Tiny (Vit-Ti/16)
18
+ """
19
+ model_args = dict(patch_size=14, embed_dim=192, depth=12, num_heads=3)
20
+ model = _create_vision_transformer('vit_tiny_patch14_224', pretrained=pretrained, **dict(model_args, **kwargs))
21
+ return model
22
+
23
+
24
+ @register_model
25
+ def vit_small_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:
26
+ """ ViT-Small (ViT-S/16)
27
+ """
28
+ model_args = dict(patch_size=14, embed_dim=384, depth=12, num_heads=6)
29
+ model = _create_vision_transformer('vit_small_patch16_224', pretrained=pretrained, **dict(model_args, **kwargs))
30
+ return model
31
+
32
+
33
+ @register_model
34
+ def vit_base_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:
35
+ """ ViT-Base (ViT-B/14) from original paper (https://arxiv.org/abs/2010.11929).
36
+ ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.
37
+ """
38
+ model_args = dict(patch_size=14, embed_dim=768, depth=12, num_heads=12)
39
+ model = _create_vision_transformer('vit_base_patch14_224', pretrained=pretrained, **dict(model_args, **kwargs))
40
+ return model
41
+
42
+
43
+ @register_model
44
+ def vit_huge_patch16_224(pretrained=False, **kwargs) -> VisionTransformer:
45
+ """ ViT-Huge model (ViT-H/16) from original paper (https://arxiv.org/abs/2010.11929).
46
+ """
47
+ model_args = dict(patch_size=16, embed_dim=1280, depth=32, num_heads=16)
48
+ if pretrained:
49
+ # There is no pretrained version of ViT-H/16, but we can adapt a ViT-H/14 for this purpose
50
+ model = _create_vision_transformer('vit_huge_patch14_clip_336', pretrained=True, **dict(model_args, pre_norm=True, **kwargs))
51
+ else:
52
+ model = _create_vision_transformer('vit_huge_patch16_224', pretrained=False, **dict(model_args, **kwargs))
53
+ return model
54
+
55
+
56
+ @register_model
57
+ def vit_huge_patch16_224_mlpnorm(pretrained=False, **kwargs) -> VisionTransformer:
58
+ """ ViT-Huge model (ViT-H/16) from original paper (https://arxiv.org/abs/2010.11929).
59
+ """
60
+ model = vit_huge_patch16_224(pretrained=pretrained, **kwargs)
61
+
62
+ for m in model.modules():
63
+ if isinstance(m, Mlp) and not isinstance(m.norm, nn.LayerNorm):
64
+ m.norm = nn.LayerNorm(m.fc1.out_features)
65
+
66
+ return model
hf_model.py ADDED
@@ -0,0 +1,174 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ from collections import namedtuple
15
+ from typing import Callable, Optional, List, Union
16
+
17
+ from timm.models import VisionTransformer
18
+ import torch
19
+ from torch import nn
20
+ from transformers import PretrainedConfig, PreTrainedModel
21
+
22
+
23
+ from .common import RESOURCE_MAP, DEFAULT_VERSION
24
+
25
+ # Import all required modules.
26
+ from .adaptor_base import AdaptorBase, RadioOutput, AdaptorInput
27
+ from .adaptor_generic import GenericAdaptor, AdaptorBase
28
+ from .adaptor_mlp import create_mlp_from_state
29
+ from .adaptor_registry import adaptor_registry
30
+ from .cls_token import ClsToken
31
+ from .enable_cpe_support import enable_cpe
32
+ from .enable_spectral_reparam import configure_spectral_reparam_from_args
33
+ from .eradio_model import eradio
34
+ from .radio_model import create_model_from_args
35
+ from .radio_model import RADIOModel as RADIOModelBase, Resolution
36
+ from .input_conditioner import get_default_conditioner, InputConditioner
37
+ from .open_clip_adaptor import OpenCLIP_RADIO
38
+ from .vit_patch_generator import ViTPatchGenerator
39
+ from .vitdet import apply_vitdet_arch, VitDetArgs
40
+
41
+ # Register extra models
42
+ from .extra_timm_models import *
43
+
44
+
45
+ class RADIOConfig(PretrainedConfig):
46
+ """Pretrained Hugging Face configuration for RADIO models."""
47
+
48
+ def __init__(
49
+ self,
50
+ args: Optional[dict] = None,
51
+ version: Optional[str] = DEFAULT_VERSION,
52
+ patch_size: Optional[int] = None,
53
+ max_resolution: Optional[int] = None,
54
+ preferred_resolution: Optional[Resolution] = None,
55
+ adaptor_names: Union[str, List[str]] = None,
56
+ vitdet_window_size: Optional[int] = None,
57
+ **kwargs,
58
+ ):
59
+ self.args = args
60
+ for field in ["dtype", "amp_dtype"]:
61
+ if self.args is not None and field in self.args:
62
+ # Convert to a string in order to make it serializable.
63
+ # For example for torch.float32 we will store "float32",
64
+ # for "bfloat16" we will store "bfloat16".
65
+ self.args[field] = str(args[field]).split(".")[-1]
66
+ self.version = version
67
+ resource = RESOURCE_MAP[version]
68
+ self.patch_size = patch_size or resource.patch_size
69
+ self.max_resolution = max_resolution or resource.max_resolution
70
+ self.preferred_resolution = (
71
+ preferred_resolution or resource.preferred_resolution
72
+ )
73
+ self.adaptor_names = adaptor_names
74
+ self.vitdet_window_size = vitdet_window_size
75
+ super().__init__(**kwargs)
76
+
77
+
78
+ class RADIOModel(PreTrainedModel):
79
+ """Pretrained Hugging Face model for RADIO.
80
+
81
+ This class inherits from PreTrainedModel, which provides
82
+ HuggingFace's functionality for loading and saving models.
83
+ """
84
+
85
+ config_class = RADIOConfig
86
+
87
+ def __init__(self, config: RADIOConfig):
88
+ super().__init__(config)
89
+
90
+ RADIOArgs = namedtuple("RADIOArgs", config.args.keys())
91
+ args = RADIOArgs(**config.args)
92
+ self.config = config
93
+
94
+ model = create_model_from_args(args)
95
+ input_conditioner: InputConditioner = get_default_conditioner()
96
+
97
+ dtype = getattr(args, "dtype", torch.float32)
98
+ if isinstance(dtype, str):
99
+ # Convert the dtype's string representation back to a dtype.
100
+ dtype = getattr(torch, dtype)
101
+ model.to(dtype=dtype)
102
+ input_conditioner.dtype = dtype
103
+
104
+ summary_idxs = torch.tensor(
105
+ [i for i, t in enumerate(args.teachers) if t.get("use_summary", True)],
106
+ dtype=torch.int64,
107
+ )
108
+
109
+ adaptor_names = config.adaptor_names
110
+ if adaptor_names is not None:
111
+ raise NotImplementedError(
112
+ f"Adaptors are not yet supported in Hugging Face models. Adaptor names: {adaptor_names}"
113
+ )
114
+
115
+ adaptors = dict()
116
+
117
+ self.radio_model = RADIOModelBase(
118
+ model,
119
+ input_conditioner,
120
+ summary_idxs=summary_idxs,
121
+ patch_size=config.patch_size,
122
+ max_resolution=config.max_resolution,
123
+ window_size=config.vitdet_window_size,
124
+ preferred_resolution=config.preferred_resolution,
125
+ adaptors=adaptors,
126
+ )
127
+
128
+ @property
129
+ def adaptors(self) -> nn.ModuleDict:
130
+ return self.radio_model.adaptors
131
+
132
+ @property
133
+ def model(self) -> VisionTransformer:
134
+ return self.radio_model.model
135
+
136
+ @property
137
+ def input_conditioner(self) -> InputConditioner:
138
+ return self.radio_model.input_conditioner
139
+
140
+ @property
141
+ def num_summary_tokens(self) -> int:
142
+ return self.radio_model.num_summary_tokens
143
+
144
+ @property
145
+ def patch_size(self) -> int:
146
+ return self.radio_model.patch_size
147
+
148
+ @property
149
+ def max_resolution(self) -> int:
150
+ return self.radio_model.max_resolution
151
+
152
+ @property
153
+ def preferred_resolution(self) -> Resolution:
154
+ return self.radio_model.preferred_resolution
155
+
156
+ @property
157
+ def window_size(self) -> int:
158
+ return self.radio_model.window_size
159
+
160
+ @property
161
+ def min_resolution_step(self) -> int:
162
+ return self.radio_model.min_resolution_step
163
+
164
+ def make_preprocessor_external(self) -> Callable[[torch.Tensor], torch.Tensor]:
165
+ return self.radio_model.make_preprocessor_external()
166
+
167
+ def get_nearest_supported_resolution(self, height: int, width: int) -> Resolution:
168
+ return self.radio_model.get_nearest_supported_resolution(height, width)
169
+
170
+ def switch_to_deploy(self):
171
+ return self.radio_model.switch_to_deploy()
172
+
173
+ def forward(self, x: torch.Tensor):
174
+ return self.radio_model.forward(x)
input_conditioner.py ADDED
@@ -0,0 +1,49 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved.
2
+ #
3
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
4
+ # and proprietary rights in and to this software, related documentation
5
+ # and any modifications thereto. Any use, reproduction, disclosure or
6
+ # distribution of this software and related documentation without an express
7
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
8
+
9
+ from typing import Union, Tuple
10
+
11
+ import torch
12
+ from torch import nn
13
+
14
+
15
+ norm_t = Union[Tuple[float, float, float], torch.Tensor]
16
+
17
+ class InputConditioner(nn.Module):
18
+ def __init__(self,
19
+ input_scale: float,
20
+ norm_mean: norm_t,
21
+ norm_std: norm_t,
22
+ dtype: torch.dtype = None,
23
+ ):
24
+ super().__init__()
25
+
26
+ self.dtype = dtype
27
+
28
+ self.register_buffer("norm_mean", _to_tensor(norm_mean) / input_scale)
29
+ self.register_buffer("norm_std", _to_tensor(norm_std) / input_scale)
30
+
31
+ def forward(self, x: torch.Tensor):
32
+ y = (x - self.norm_mean) / self.norm_std
33
+ if self.dtype is not None:
34
+ y = y.to(self.dtype)
35
+ return y
36
+
37
+
38
+ def get_default_conditioner():
39
+ from timm.data.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD
40
+
41
+ return InputConditioner(
42
+ input_scale=1.0,
43
+ norm_mean=OPENAI_CLIP_MEAN,
44
+ norm_std=OPENAI_CLIP_STD,
45
+ )
46
+
47
+
48
+ def _to_tensor(v: norm_t):
49
+ return torch.as_tensor(v, dtype=torch.float32).view(-1, 1, 1)
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:608c1132b7a82bf4652827a209509864b5d7d925aa2d2d256e2e1a2968b16792
3
+ size 392950088
open_clip_adaptor.py ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
2
+ #
3
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
4
+ # and proprietary rights in and to this software, related documentation
5
+ # and any modifications thereto. Any use, reproduction, disclosure or
6
+ # distribution of this software and related documentation without an express
7
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
8
+ from argparse import Namespace
9
+
10
+ import torch
11
+ from torch import nn
12
+ import torch.nn.functional as F
13
+
14
+ from .adaptor_registry import adaptor_registry, dict_t, state_t
15
+
16
+ from .adaptor_generic import GenericAdaptor
17
+
18
+
19
+ class OpenCLIP_RADIO(GenericAdaptor):
20
+ def __init__(self, main_config: Namespace, adaptor_config: dict_t, state: state_t):
21
+ super().__init__(main_config, adaptor_config, state)
22
+
23
+ import open_clip
24
+
25
+ self.oc_model = open_clip.create_model_from_pretrained(
26
+ model_name=adaptor_config['model'],
27
+ pretrained=adaptor_config['pretrained'],
28
+ return_transform=False,
29
+ )
30
+ # Unload these parameters
31
+ self.oc_model.visual = None
32
+
33
+ self.tokenizer = open_clip.get_tokenizer(model_name=adaptor_config['model'])
34
+
35
+ def encode_text(self, text, normalize: bool = False):
36
+ return self.oc_model.encode_text(text, normalize=normalize)
37
+
38
+
39
+ @adaptor_registry.register_adaptor("open_clip")
40
+ def create_open_clip_adaptor(main_config: Namespace, adaptor_config: dict_t, state: state_t):
41
+ return OpenCLIP_RADIO(main_config, adaptor_config, state)
radio_model.py ADDED
@@ -0,0 +1,257 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved.
2
+ #
3
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
4
+ # and proprietary rights in and to this software, related documentation
5
+ # and any modifications thereto. Any use, reproduction, disclosure or
6
+ # distribution of this software and related documentation without an express
7
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
8
+ from typing import Callable, Dict, List, NamedTuple, Optional, Tuple, Union
9
+
10
+ import torch
11
+ from torch import nn
12
+
13
+ from timm.models import create_model, VisionTransformer
14
+
15
+ from .enable_cpe_support import enable_cpe
16
+ from .input_conditioner import InputConditioner
17
+ # Register extra models
18
+ from . import extra_timm_models
19
+ from .adaptor_base import AdaptorBase, RadioOutput, AdaptorInput
20
+ from . import eradio_model
21
+ from .enable_spectral_reparam import configure_spectral_reparam_from_args
22
+
23
+
24
+ class Resolution(NamedTuple):
25
+ height: int
26
+ width: int
27
+
28
+
29
+ class RADIOModel(nn.Module):
30
+ def __init__(
31
+ self,
32
+ model: nn.Module,
33
+ input_conditioner: InputConditioner,
34
+ patch_size: int,
35
+ max_resolution: int,
36
+ preferred_resolution: Resolution,
37
+ summary_idxs: Optional[torch.Tensor] = None,
38
+ window_size: int = None,
39
+ adaptors: Dict[str, AdaptorBase] = None,
40
+ ):
41
+ super().__init__()
42
+
43
+ self.model = model
44
+ self.input_conditioner = input_conditioner
45
+ if summary_idxs is not None:
46
+ self.register_buffer('summary_idxs', summary_idxs)
47
+ else:
48
+ self.summary_idxs = None
49
+
50
+ self._preferred_resolution = preferred_resolution
51
+ self._patch_size = patch_size
52
+ self._max_resolution = max_resolution
53
+ self._window_size = window_size
54
+
55
+ adaptors = adaptors or dict()
56
+ self.adaptors = nn.ModuleDict(adaptors)
57
+
58
+ @property
59
+ def num_summary_tokens(self) -> int:
60
+ patch_gen = getattr(self.model, "patch_generator", None)
61
+ if patch_gen is not None:
62
+ return patch_gen.num_skip
63
+ elif self.model.global_pool == 'avg':
64
+ return 0
65
+ return 1
66
+
67
+ @property
68
+ def patch_size(self) -> int:
69
+ if self._patch_size is not None:
70
+ return self._patch_size
71
+ patch_gen = getattr(self.model, "patch_generator", None)
72
+ if patch_gen is not None:
73
+ return patch_gen.patch_size
74
+ return None
75
+
76
+ @property
77
+ def max_resolution(self) -> int:
78
+ return self._max_resolution
79
+
80
+ @property
81
+ def preferred_resolution(self) -> Resolution:
82
+ return self._preferred_resolution
83
+
84
+ @property
85
+ def window_size(self) -> int:
86
+ return self._window_size
87
+
88
+ @property
89
+ def min_resolution_step(self) -> int:
90
+ res = self.patch_size
91
+ if self.window_size is not None:
92
+ res *= self.window_size
93
+ return res
94
+
95
+ def make_preprocessor_external(self) -> Callable[[torch.Tensor], torch.Tensor]:
96
+ ret = self.input_conditioner
97
+ self.input_conditioner = nn.Identity()
98
+ return ret
99
+
100
+ def get_nearest_supported_resolution(self, height: int, width: int) -> Resolution:
101
+ height = int(round(height / self.min_resolution_step) * self.min_resolution_step)
102
+ width = int(round(width / self.min_resolution_step) * self.min_resolution_step)
103
+
104
+ height = max(height, self.min_resolution_step)
105
+ width = max(width, self.min_resolution_step)
106
+
107
+ return Resolution(height=height, width=width)
108
+
109
+ def switch_to_deploy(self):
110
+ fn = getattr(self.model, 'switch_to_deploy', None)
111
+ if fn is not None:
112
+ fn()
113
+
114
+ def forward(self, x: torch.Tensor) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
115
+ res_step = self.min_resolution_step
116
+ if res_step is not None and (x.shape[-2] % res_step != 0 or x.shape[-1] % res_step != 0):
117
+ raise ValueError('The input resolution must be a multiple of `self.min_resolution_step`. '
118
+ '`self.get_nearest_supported_resolution(<height>, <width>) is provided as a convenience API. '
119
+ f'Input: {x.shape[-2:]}, Nearest: {self.get_nearest_supported_resolution(*x.shape[-2:])}')
120
+
121
+ x = self.input_conditioner(x)
122
+ y = self.model.forward_features(x)
123
+
124
+ if isinstance(self.model, VisionTransformer):
125
+ patch_gen = getattr(self.model, "patch_generator", None)
126
+ if patch_gen is not None:
127
+ all_summary = y[:, : patch_gen.num_cls_tokens]
128
+ if self.summary_idxs is not None:
129
+ bb_summary = all_summary[:, self.summary_idxs]
130
+ else:
131
+ bb_summary = all_summary
132
+ all_feat = y[:, patch_gen.num_skip :]
133
+ elif self.model.global_pool == "avg":
134
+ all_summary = y[:, self.model.num_prefix_tokens :].mean(dim=1)
135
+ bb_summary = all_summary
136
+ all_feat = y
137
+ else:
138
+ all_summary = y[:, 0]
139
+ bb_summary = all_summary
140
+ all_feat = y[:, 1:]
141
+ elif isinstance(self.model, eradio_model.ERADIO):
142
+ _, f = y
143
+ all_feat = f.flatten(2).transpose(1, 2)
144
+ all_summary = all_feat.mean(dim=1)
145
+ bb_summary = all_summary
146
+ elif isinstance(y, (list, tuple)):
147
+ all_summary, all_feat = y
148
+ bb_summary = all_summary
149
+ else:
150
+ raise ValueError("Unsupported model type")
151
+
152
+ all_feat = all_feat.float()
153
+ ret = RadioOutput(bb_summary.flatten(1), all_feat).to(torch.float32)
154
+ if self.adaptors:
155
+ ret = dict(backbone=ret)
156
+ for name, adaptor in self.adaptors.items():
157
+ if all_summary.ndim == 3:
158
+ summary = all_summary[:, adaptor.head_idx]
159
+ else:
160
+ summary = all_summary
161
+ ada_input = AdaptorInput(images=x, summary=summary.float(), features=all_feat)
162
+ v = adaptor(ada_input).to(torch.float32)
163
+ ret[name] = v
164
+
165
+ return ret
166
+
167
+ def forward_intermediates(
168
+ self,
169
+ x: torch.Tensor,
170
+ indices: Optional[Union[int, List[int], Tuple[int]]] = None,
171
+ return_prefix_tokens: bool = False,
172
+ norm: bool = False,
173
+ stop_early: bool = False,
174
+ output_fmt: str = 'NCHW',
175
+ intermediates_only: bool = False,
176
+ aggregation: Optional[str] = "sparse",
177
+ ) -> List[RadioOutput]:
178
+ """ Forward features that returns intermediates.
179
+ Args:
180
+ x: Input image tensor
181
+ indices: Take last n blocks if int, select matching indices if sequence
182
+ return_prefix_tokens: Return both prefix and spatial intermediate tokens
183
+ norm: Apply norm layer to all intermediates
184
+ stop_early: Stop iterating over blocks when last desired intermediate hit
185
+ output_fmt: Shape of intermediate feature outputs
186
+ intermediates_only: Only return intermediate features
187
+ aggregation: intermediate layer aggregation method (sparse or dense).
188
+ Dense accumulation is done by averaging the features in each group.
189
+ Returns:
190
+ List of RadioOutput objects.
191
+ """
192
+ outputs = self.model.forward_intermediates(
193
+ x,
194
+ indices=indices,
195
+ return_prefix_tokens=return_prefix_tokens,
196
+ norm=norm,
197
+ stop_early=stop_early,
198
+ output_fmt=output_fmt,
199
+ intermediates_only=intermediates_only,
200
+ aggregation=aggregation,
201
+ )
202
+ if return_prefix_tokens:
203
+ radio_outputs = [RadioOutput(summary, features) for (summary, features) in outputs]
204
+ else:
205
+ radio_outputs = [RadioOutput(None, features) for features in outputs]
206
+ return radio_outputs
207
+
208
+
209
+ def create_model_from_args(args) -> nn.Module:
210
+ in_chans = 3
211
+ if args.in_chans is not None:
212
+ in_chans = args.in_chans
213
+ elif args.input_size is not None:
214
+ in_chans = args.input_size[0]
215
+
216
+ # Skip weight initialization unless it's explicitly requested.
217
+ weight_init = args.model_kwargs.pop("weight_init", "skip")
218
+
219
+ model = create_model(
220
+ args.model,
221
+ pretrained=args.pretrained,
222
+ in_chans=in_chans,
223
+ num_classes=args.num_classes,
224
+ drop_rate=args.drop,
225
+ drop_path_rate=args.drop_path,
226
+ drop_block_rate=args.drop_block,
227
+ global_pool=args.gp,
228
+ bn_momentum=args.bn_momentum,
229
+ bn_eps=args.bn_eps,
230
+ scriptable=args.torchscript,
231
+ checkpoint_path=args.initial_checkpoint,
232
+ weight_init=weight_init,
233
+ **args.model_kwargs,
234
+ )
235
+
236
+ if hasattr(model, 'norm') and not getattr(args, 'model_norm', False):
237
+ model.norm = nn.Identity()
238
+
239
+ model.head = nn.Identity()
240
+
241
+ assert (
242
+ not args.cls_token_per_teacher or args.cpe_max_size is not None
243
+ ), "CPE must be enabled for multiple CLS tokens!"
244
+
245
+ if args.cpe_max_size is not None:
246
+ uq_teachers = set(t['name'] for t in args.teachers)
247
+ enable_cpe(
248
+ model,
249
+ args.cpe_max_size,
250
+ num_cls_tokens=len(uq_teachers) if args.cls_token_per_teacher else 1,
251
+ register_multiple=args.register_multiple,
252
+ )
253
+
254
+ if args.spectral_reparam:
255
+ configure_spectral_reparam_from_args(model, args)
256
+
257
+ return model
vit_patch_generator.py ADDED
@@ -0,0 +1,299 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved.
2
+ #
3
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
4
+ # and proprietary rights in and to this software, related documentation
5
+ # and any modifications thereto. Any use, reproduction, disclosure or
6
+ # distribution of this software and related documentation without an express
7
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
8
+
9
+ import math
10
+ from typing import Union, Tuple, Optional
11
+
12
+ import torch
13
+ import torch.nn.functional as F
14
+ from torch import nn
15
+ from einops import rearrange
16
+
17
+ from .cls_token import ClsToken
18
+
19
+ input_dim_t = Union[int, Tuple[int, int]]
20
+
21
+ try:
22
+ # raise ImportError()
23
+ from indirect_grid_sample import indirect_grid_sample
24
+ except ImportError:
25
+ indirect_grid_sample = None
26
+
27
+ class ViTPatchGenerator(nn.Module):
28
+ def __init__(self,
29
+ patch_size: int,
30
+ embed_dim: int,
31
+ input_dims: input_dim_t,
32
+ abs_pos: bool = True,
33
+ normalize_patches: bool = False,
34
+ cls_token: bool = False,
35
+ max_input_dims: Optional[input_dim_t] = None,
36
+ pos_dropout: float = 0.0,
37
+ return_pos_enc: bool = False,
38
+ num_cls_tokens: int = 1,
39
+ register_multiple: int = 0,
40
+ device=None, dtype=None,
41
+ ):
42
+ super().__init__()
43
+
44
+ if isinstance(input_dims, int):
45
+ input_dims = (input_dims, input_dims)
46
+
47
+ if max_input_dims is None:
48
+ max_input_dims = input_dims
49
+ if isinstance(max_input_dims, int):
50
+ max_input_dims = (max_input_dims, max_input_dims)
51
+
52
+ max_input_dims = tuple(
53
+ int(math.ceil(d / patch_size) * patch_size)
54
+ for d in max_input_dims
55
+ )
56
+
57
+ self.cpe_mode = max_input_dims != input_dims
58
+ self.pos_dropout = pos_dropout
59
+ self.return_pos_enc = return_pos_enc
60
+
61
+ factory = dict(device=device, dtype=dtype)
62
+
63
+ self.patch_size = patch_size
64
+ self.abs_pos = abs_pos
65
+ self.embed_dim = embed_dim
66
+
67
+ self.num_rows = max_input_dims[0] // patch_size
68
+ self.num_cols = max_input_dims[1] // patch_size
69
+ self.input_dims = tuple(d // patch_size for d in input_dims)
70
+ self.num_patches = self.num_rows * self.num_cols
71
+ self.max_input_dims = max_input_dims
72
+
73
+ self.im_to_patches = Im2Patches(patch_size)
74
+ self.embedder = ViTPatchLinear(patch_size, embed_dim, **factory)
75
+
76
+ if abs_pos:
77
+ scale = embed_dim ** -0.5
78
+ self.pos_embed = nn.Parameter(torch.randn(1, self.num_patches, embed_dim, **factory) * scale)
79
+
80
+ self.cls_token = ClsToken(
81
+ embed_dim,
82
+ num_tokens=num_cls_tokens,
83
+ enabled=cls_token,
84
+ register_multiple=register_multiple,
85
+ )
86
+
87
+ self.patch_normalizer = nn.LayerNorm(embed_dim) if normalize_patches else nn.Identity()
88
+
89
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
90
+ patches = self.embed_patches(x)
91
+ patches, pos_enc = self.apply_pos_enc(patches, input_size=x.shape[2:])
92
+ patches = self.cls_token(patches)
93
+ patches = self.patch_normalizer(patches)
94
+ if self.return_pos_enc:
95
+ return patches, pos_enc
96
+ return patches
97
+
98
+ @property
99
+ def apply_cls_token(self):
100
+ return self.cls_token.enabled
101
+
102
+ @property
103
+ def num_cls_tokens(self):
104
+ return self.cls_token.num_tokens
105
+
106
+ @property
107
+ def num_registers(self):
108
+ return self.cls_token.num_registers
109
+
110
+ @property
111
+ def num_skip(self):
112
+ return self.num_cls_tokens + self.num_registers
113
+
114
+ def no_weight_decay(self):
115
+ return [
116
+ 'pos_embed',
117
+ ]
118
+
119
+ def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
120
+ if self.abs_pos:
121
+ self._load_embed(state_dict[f'{prefix}pos_embed'], self.pos_embed)
122
+
123
+ def _load_embed(self, src_embed: torch.Tensor, targ_embed: nn.Parameter):
124
+ if src_embed.shape != targ_embed.shape:
125
+ src_size = int(math.sqrt(src_embed.shape[1]))
126
+
127
+ assert src_size ** 2 == src_embed.shape[1], 'Unable to interpolate non-square embedding'
128
+
129
+ src_embed = rearrange(src_embed, 'b (h w) c -> b c h w', h=src_size, w=src_size)
130
+ src_embed = F.interpolate(src_embed, size=(self.num_rows, self.num_cols), mode='bicubic', align_corners=True, antialias=False)
131
+ src_embed = rearrange(src_embed, 'b c h w -> b (h w) c')
132
+ targ_embed.data.copy_(src_embed)
133
+
134
+ def _load_projection(self, src_proj_weight: torch.Tensor, targ_proj_weight: torch.Tensor):
135
+ if src_proj_weight.shape != targ_proj_weight.shape:
136
+ src_patch_size = int(math.sqrt(src_proj_weight.shape[1] // 3))
137
+
138
+ assert (src_patch_size ** 2) * 3 == src_proj_weight.shape[1], 'Unable to interpolate non-square patch size'
139
+
140
+ src_proj_weight = rearrange(src_proj_weight, 'b (c h w) -> b c h w', c=3, h=src_patch_size, w=src_patch_size)
141
+ src_proj_weight = F.interpolate(src_proj_weight, size=(self.patch_size, self.patch_size), mode='bicubic', align_corners=True, antialias=False)
142
+ src_proj_weight = rearrange(src_proj_weight, 'b c h w -> b (c h w)')
143
+ targ_proj_weight.data.copy_(src_proj_weight)
144
+
145
+ def embed_patches(self, x: torch.Tensor) -> torch.Tensor:
146
+ patches = self.im_to_patches(x)
147
+ patches = self.embedder(patches)
148
+ return patches
149
+
150
+ def apply_pos_enc(self,
151
+ patches: torch.Tensor,
152
+ patch_idxs: Optional[torch.Tensor] = None,
153
+ input_size: Optional[Tuple[int, int]] = None,
154
+ ) -> torch.Tensor:
155
+ if not self.abs_pos:
156
+ return patches
157
+
158
+ pos_enc = self.get_pos_enc(patches.shape[0], patch_idxs, input_size)
159
+
160
+ if self.training and self.pos_dropout > 0:
161
+ keeps = torch.rand(patches.shape[0], 1, 1, dtype=pos_enc.dtype, device=pos_enc.device) > self.pos_dropout
162
+ pos_enc_drop = torch.where(keeps, pos_enc, 0)
163
+ else:
164
+ pos_enc_drop = pos_enc
165
+
166
+ return patches + pos_enc_drop, pos_enc
167
+
168
+ def get_pos_enc(self,
169
+ batch_size: int,
170
+ patch_idxs: Optional[torch.Tensor] = None,
171
+ input_size: Optional[Tuple[int, int]] = None,
172
+ ) -> torch.Tensor:
173
+ if input_size is None:
174
+ input_dims = self.input_dims
175
+ else:
176
+ input_dims = tuple(d // self.patch_size for d in input_size)
177
+
178
+ pos_embed = self._get_pos_embeddings(batch_size, input_dims)
179
+
180
+ if patch_idxs is None:
181
+ return pos_embed
182
+
183
+ exp_patch_idxs = patch_idxs.unsqueeze(-1).expand(-1, -1, pos_embed.shape[-1])
184
+
185
+ pos_embed = torch.gather(pos_embed.expand(patch_idxs.shape[0], -1, -1), dim=1, index=exp_patch_idxs)
186
+ return pos_embed
187
+
188
+
189
+ def _get_pos_embeddings(self, batch_size: int, input_dims: Tuple[int, int]):
190
+ if (self.num_rows, self.num_cols) == input_dims:
191
+ return self.pos_embed
192
+
193
+ pos_embed = self.pos_embed.reshape(1, self.num_rows, self.num_cols, -1).permute(0, 3, 1, 2)
194
+
195
+ def window_select(pos_embed):
196
+ if input_dims[0] < pos_embed.shape[-2]:
197
+ pos_embed = pos_embed[..., :input_dims[0], :]
198
+ if input_dims[1] < pos_embed.shape[-1]:
199
+ pos_embed = pos_embed[..., :, :input_dims[1]]
200
+ return pos_embed
201
+
202
+ if self.cpe_mode:
203
+ if self.training:
204
+ min_scale = math.sqrt(0.1)
205
+ scale = torch.rand(batch_size, 1, 1, device=pos_embed.device) * (1 - min_scale) + min_scale
206
+ aspect_min = math.log(3 / 4)
207
+ aspect_max = -aspect_min
208
+ aspect = torch.exp(torch.rand(batch_size, 1, 1, device=pos_embed.device) * (aspect_max - aspect_min) + aspect_min)
209
+
210
+ scale_x = scale * aspect
211
+ scale_y = scale * (1 / aspect)
212
+ scale_xy = torch.stack([scale_x, scale_y], dim=-1).clamp_(0, 1)
213
+
214
+ pos_xy = torch.rand(batch_size, 1, 1, 2, device=pos_embed.device) * (1 - scale_xy)
215
+
216
+ lin_x = torch.linspace(0, 1, steps=input_dims[1], device=pos_embed.device)[None, None].expand(batch_size, input_dims[0], -1)
217
+ lin_y = torch.linspace(0, 1, steps=input_dims[0], device=pos_embed.device)[None, :, None].expand(batch_size, -1, input_dims[1])
218
+
219
+ lin_xy = torch.stack([lin_x, lin_y], dim=-1)
220
+
221
+ grid_xy = lin_xy * scale_xy + pos_xy
222
+
223
+ # Convert to [-1, 1] range
224
+ grid_xy.mul_(2).sub_(1)
225
+
226
+ pos_embed = F.grid_sample(
227
+ pos_embed.float().expand(batch_size, -1, -1, -1),
228
+ grid=grid_xy,
229
+ mode='bilinear',
230
+ padding_mode='zeros',
231
+ align_corners=True,
232
+ ).to(pos_embed.dtype)
233
+ else:
234
+ # i_rows, i_cols = input_dims
235
+ # p_rows, p_cols = pos_embed.shape[2:]
236
+ # if i_rows <= p_rows and i_cols <= p_cols:
237
+ # left = (p_cols - i_cols) // 2
238
+ # top = (p_rows - i_rows) // 2
239
+ # pos_embed = pos_embed[..., top:top+i_rows, left:left+i_cols]
240
+ # else:
241
+ max_dim = max(input_dims)
242
+ pos_embed = F.interpolate(pos_embed.float(), size=(max_dim, max_dim), align_corners=True, mode='bilinear').to(pos_embed.dtype)
243
+
244
+ pos_embed = window_select(pos_embed)
245
+ else:
246
+ pos_embed = window_select(pos_embed)
247
+
248
+ if pos_embed.shape[-2:] != input_dims:
249
+ pos_embed = F.interpolate(pos_embed.float(), size=input_dims, align_corners=True, mode='bilinear').to(pos_embed.dtype)
250
+
251
+ pos_embed = pos_embed.flatten(2).permute(0, 2, 1)
252
+
253
+ return pos_embed
254
+
255
+
256
+ class Im2Patches(nn.Module):
257
+ def __init__(self, patch_size: int):
258
+ super().__init__()
259
+ self.patch_size = patch_size
260
+
261
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
262
+ if self.patch_size == 1:
263
+ patches = x.flatten(2)
264
+ patches = patches.permute(0, 2, 1)
265
+ return patches
266
+
267
+ py = x.shape[-2] // self.patch_size
268
+ px = x.shape[-1] // self.patch_size
269
+ patches = rearrange(x, 'b c (py yy) (px xx) -> b (py px) (c yy xx)',
270
+ py=py, yy=self.patch_size,
271
+ px=px, xx=self.patch_size,
272
+ )
273
+ return patches
274
+
275
+
276
+ class ViTPatchLinear(nn.Linear):
277
+ def __init__(self, patch_size: int, embed_dim: int, **factory):
278
+ super().__init__(
279
+ 3 * (patch_size ** 2),
280
+ embed_dim,
281
+ bias=False,
282
+ **factory
283
+ )
284
+ self.patch_size = patch_size
285
+
286
+ def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
287
+ if self.bias is not None:
288
+ self.bias.data.copy_(state_dict[f'{prefix}bias'])
289
+
290
+ chk_weight = state_dict[f'{prefix}weight']
291
+ if chk_weight.shape != self.weight.shape:
292
+ src_patch_size = int(math.sqrt(chk_weight.shape[1] // 3))
293
+
294
+ assert (src_patch_size ** 2) * 3 == chk_weight.shape[1], 'Unable to interpolate non-square patch size'
295
+
296
+ chk_weight = rearrange(chk_weight, 'b (c h w) -> b c h w', c=3, h=src_patch_size, w=src_patch_size)
297
+ chk_weight = F.interpolate(chk_weight, size=(self.patch_size, self.patch_size), mode='bicubic', align_corners=True, antialias=False)
298
+ chk_weight = rearrange(chk_weight, 'b c h w -> b (c h w)')
299
+ self.weight.data.copy_(chk_weight)
vitdet.py ADDED
@@ -0,0 +1,181 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from collections import defaultdict
2
+ from contextlib import contextmanager
3
+ from logging import getLogger
4
+ import math
5
+ import sys
6
+ from typing import List, Union, Iterable
7
+
8
+ import numpy as np
9
+ import torch
10
+ from torch import nn
11
+
12
+ from timm.models import VisionTransformer
13
+ from einops import rearrange
14
+
15
+ DEFAULT_NUM_WINDOWED = 5
16
+ DEFAULT_NUM_GLOBAL = 4
17
+
18
+
19
+ class VitDetArgs:
20
+ def __init__(self,
21
+ window_size: int,
22
+ num_summary_tokens: int,
23
+ num_windowed: int = None,
24
+ num_global: int = None,
25
+ ):
26
+ self.window_size = window_size
27
+ self.num_summary_tokens = num_summary_tokens
28
+ self.num_windowed = num_windowed
29
+ self.num_global = num_global
30
+
31
+
32
+ def apply_vitdet_arch(model: VisionTransformer, args: VitDetArgs):
33
+ if isinstance(model, VisionTransformer):
34
+ patch_embed = getattr(model, 'patch_generator', model.patch_embed)
35
+
36
+ return ViTDetHook(patch_embed, model.blocks, args)
37
+ else:
38
+ print(f'Warning: Unable to apply VitDet aug!', file=sys.stderr)
39
+
40
+
41
+ class ViTDetHook:
42
+ def __init__(self,
43
+ embedder: nn.Module,
44
+ blocks: nn.Sequential,
45
+ args: VitDetArgs,
46
+ ):
47
+ self.blocks = blocks
48
+ self.num_summary_tokens = args.num_summary_tokens
49
+ self.window_size = args.window_size
50
+
51
+ self._input_resolution = None
52
+ self._num_windows = None
53
+ self._cls_patch = None
54
+ self._order_cache = dict()
55
+
56
+ embedder.register_forward_pre_hook(self._enter_model)
57
+
58
+ # This will decide if we window-fy the patches
59
+ # and enable vit-det for this iteration, and if so,
60
+ # rearrange the patches for efficient mode switching
61
+ blocks.register_forward_pre_hook(self._enter_blocks)
62
+
63
+ is_global = True
64
+ if args.num_windowed is not None:
65
+ period = args.num_windowed + 1
66
+ else:
67
+ num_global = args.num_global or DEFAULT_NUM_GLOBAL
68
+ period = max(len(blocks) // num_global, 1)
69
+
70
+ for i, layer in enumerate(blocks[:-1]):
71
+ ctr = i % period
72
+ if ctr == 0:
73
+ layer.register_forward_pre_hook(self._to_windows)
74
+ is_global = False
75
+ elif ctr == period - 1:
76
+ layer.register_forward_pre_hook(self._to_global)
77
+ is_global = True
78
+
79
+ # Always ensure the final layer is a global layer
80
+ if not is_global:
81
+ blocks[-1].register_forward_pre_hook(self._to_global)
82
+
83
+ blocks.register_forward_hook(self._exit_model)
84
+
85
+ def _enter_model(self, _, input: List[torch.Tensor]):
86
+ self._input_resolution = input[0].shape[-2:]
87
+
88
+ def _enter_blocks(self, _, input: List[torch.Tensor]):
89
+ # print(f'{get_rank()} - ViTDet Window Size: {self._window_size}', file=sys.stderr)
90
+
91
+ patches = input[0]
92
+ patches = self._rearrange_patches(patches)
93
+
94
+ return (patches,) + input[1:]
95
+
96
+ def _to_windows(self, _, input: List[torch.Tensor]):
97
+ patches = input[0]
98
+
99
+ if self.num_summary_tokens:
100
+ self._cls_patch = patches[:, :self.num_summary_tokens]
101
+ patches = patches[:, self.num_summary_tokens:]
102
+
103
+ patches = rearrange(
104
+ patches, 'b (p t) c -> (b p) t c',
105
+ p=self._num_windows, t=self.window_size ** 2,
106
+ )
107
+
108
+ return (patches,) + input[1:]
109
+
110
+ def _to_global(self, _, input: List[torch.Tensor]):
111
+ patches = input[0]
112
+
113
+ patches = rearrange(
114
+ patches, '(b p) t c -> b (p t) c',
115
+ p=self._num_windows, t=self.window_size ** 2,
116
+ b=patches.shape[0] // self._num_windows,
117
+ )
118
+
119
+ if self.num_summary_tokens:
120
+ patches = torch.cat([
121
+ self._cls_patch,
122
+ patches,
123
+ ], dim=1)
124
+
125
+ return (patches,) + input[1:]
126
+
127
+ def _exit_model(self, _, inputs: List[torch.Tensor], patches: torch.Tensor):
128
+ # Return patches to their original order
129
+ patch_order = self._order_cache[self._input_resolution][0]
130
+ patch_order = patch_order.reshape(1, -1, 1).expand_as(patches)
131
+
132
+ ret_patches = torch.empty_like(patches)
133
+ ret_patches = torch.scatter(
134
+ ret_patches,
135
+ dim=1,
136
+ index=patch_order,
137
+ src=patches,
138
+ )
139
+
140
+ return ret_patches
141
+
142
+ def _rearrange_patches(self, patches: torch.Tensor):
143
+ # We rearrange the patches so that we can efficiently
144
+ # switch between windowed and global mode by just
145
+ # reshaping the tensor
146
+
147
+ patch_order, self._num_windows = self._order_cache.get(self._input_resolution, (None, None))
148
+ if patch_order is None:
149
+ num_feat_patches = patches.shape[1] - self.num_summary_tokens
150
+ num_pixels = self._input_resolution[0] * self._input_resolution[1]
151
+
152
+ patch_size = int(round(math.sqrt(num_pixels / num_feat_patches)))
153
+ rows = self._input_resolution[-2] // patch_size
154
+ cols = self._input_resolution[-1] // patch_size
155
+
156
+ w_rows = rows // self.window_size
157
+ w_cols = cols // self.window_size
158
+
159
+ patch_order = torch.arange(0, num_feat_patches, device=patches.device)
160
+
161
+ patch_order = rearrange(
162
+ patch_order, '(wy py wx px) -> (wy wx py px)',
163
+ wy=w_rows, wx=w_cols,
164
+ py=self.window_size, px=self.window_size,
165
+ )
166
+
167
+ if self.num_summary_tokens:
168
+ patch_order = torch.cat([
169
+ torch.arange(self.num_summary_tokens, dtype=patch_order.dtype, device=patch_order.device),
170
+ patch_order + self.num_summary_tokens,
171
+ ])
172
+
173
+ self._num_windows = w_rows * w_cols
174
+ self._order_cache[self._input_resolution] = (
175
+ patch_order,
176
+ self._num_windows,
177
+ )
178
+
179
+ patch_order = patch_order.reshape(1, -1, 1).expand_as(patches)
180
+ patches = torch.gather(patches, dim=1, index=patch_order)
181
+ return patches