adaptor_base.py CHANGED
@@ -6,7 +6,7 @@
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
@@ -17,6 +17,8 @@ class AdaptorInput(NamedTuple):
17
  images: torch.Tensor
18
  summary: torch.Tensor
19
  features: torch.Tensor
 
 
20
 
21
 
22
  class RadioOutput(NamedTuple):
 
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, Optional
10
 
11
  import torch
12
  from torch import nn
 
17
  images: torch.Tensor
18
  summary: torch.Tensor
19
  features: torch.Tensor
20
+ feature_fmt: str
21
+ patch_size: int
22
 
23
 
24
  class RadioOutput(NamedTuple):
adaptor_generic.py CHANGED
@@ -12,18 +12,58 @@ 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)
 
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, create_mlp_from_config
16
 
17
 
18
  class GenericAdaptor(AdaptorBase):
19
+ def __init__(self, main_config: Namespace, adaptor_config, state, mlp_config=None):
20
  super().__init__()
21
 
22
+ extra_args = dict()
23
+ ups = None
24
+ ups_rank = None
25
+ if adaptor_config is not None:
26
+ ups = adaptor_config.get('fd_upsample_factor', None)
27
+ ups_rank = adaptor_config.get('fd_upsample_rank', None)
28
+ elif mlp_config is not None:
29
+ ups = mlp_config["feature"].get('upsample_factor', None)
30
+ ups_rank = mlp_config["feature"].get('upsample_rank', None)
31
+ if ups is not None:
32
+ extra_args['upsample_factor'] = ups
33
+ extra_args['upsample_rank'] = ups_rank
34
+
35
+ if state is not None:
36
+ spectral_heads = getattr(main_config, 'spectral_heads', False)
37
+ self.head_mlp = create_mlp_from_state(main_config.mlp_version, state, 'summary.', spectral_weights=spectral_heads)
38
+ self.feat_mlp = create_mlp_from_state(main_config.mlp_version, state, 'feature.', spectral_weights=spectral_heads, **extra_args)
39
+ else:
40
+ assert mlp_config is not None, "Config must not be None if state is None"
41
+
42
+ self.head_mlp = create_mlp_from_config(
43
+ main_config.mlp_version,
44
+ mlp_config["summary"]["input_dim"],
45
+ mlp_config["summary"]["hidden_dim"],
46
+ mlp_config["summary"]["output_dim"],
47
+ mlp_config["summary"]["num_inner"],
48
+ )
49
+ self.feat_mlp = create_mlp_from_config(
50
+ main_config.mlp_version,
51
+ mlp_config["feature"]["input_dim"],
52
+ mlp_config["feature"]["hidden_dim"],
53
+ mlp_config["feature"]["output_dim"],
54
+ mlp_config["feature"]["num_inner"],
55
+ **extra_args
56
+ )
57
 
58
  def forward(self, input: AdaptorInput) -> RadioOutput:
59
+ # Convert input'd type to the type of the first parameter of the adaptor.
60
+ first_param = next(self.parameters())
61
+ summary = self.head_mlp(input.summary.to(dtype=first_param.dtype)).to(dtype=input.summary.dtype)
62
+ feat = self.feat_mlp(input.features.to(dtype=first_param.dtype), images=input.images, patch_size=input.patch_size).to(dtype=input.features.dtype)
63
+
64
+ if input.feature_fmt == 'NCHW':
65
+ feat = (feat.reshape(feat.shape[0], input.images.shape[-2] // input.patch_size * self.feat_mlp.upsample_factor, input.images.shape[-1] // input.patch_size * self.feat_mlp.upsample_factor, feat.shape[2])
66
+ .permute(0, 3, 1, 2)
67
+ )
68
 
69
  return RadioOutput(summary, feat)
adaptor_mlp.py CHANGED
@@ -6,7 +6,7 @@
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
@@ -14,6 +14,8 @@ from torch import nn
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,
@@ -51,6 +53,8 @@ class MLP2(nn.Module):
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
 
@@ -60,10 +64,12 @@ class MLP2(nn.Module):
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
 
@@ -82,7 +88,7 @@ class MLP2(nn.Module):
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:
@@ -90,8 +96,12 @@ class MLP2(nn.Module):
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
 
@@ -113,20 +123,22 @@ def strip_prefix(state: Dict[str, torch.Tensor], prefix: str):
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'
@@ -138,13 +150,25 @@ def get_mlp_info_from_state(version: str, state: Dict[str, torch.Tensor], prefix
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
 
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, Optional
10
 
11
  import torch
12
  from torch import nn
 
14
  from einops import rearrange
15
  from timm.models.vision_transformer import Block
16
 
17
+ from .enable_spectral_reparam import disable_spectral_reparam, enable_spectral_reparam
18
+
19
 
20
  class MLP(nn.Module):
21
  def __init__(self, input_size: int, hidden_size: int, output_size: int,
 
53
  num_inner: int = 0,
54
  pre_norm: bool = False, device: torch.device = None,
55
  upsample_factor: int = 1,
56
+ upsample_rank: int = None,
57
+ from_config: bool = False,
58
  **kwargs):
59
  super().__init__()
60
 
 
64
  ) if pre_norm else nn.Identity()
65
 
66
  self.upsample_factor = upsample_factor
67
+ sq_ups = upsample_factor ** 2
68
+
69
+ self._real_output_dim = output_size // sq_ups
70
 
71
+ # hidden_size *= upsample_factor
72
+ # output_size *= (upsample_factor ** 2)
73
 
74
  self.fc1 = nn.Linear(input_size, hidden_size, device=device)
75
 
 
88
  nn.Linear(hidden_size, output_size, device=device),
89
  )
90
 
91
+ def forward(self, x: torch.Tensor, images: Optional[torch.Tensor] = None, patch_size: Optional[int] = None) -> torch.Tensor:
92
  x = self.pre_norm(x)
93
  x = self.fc1(x)
94
  for block in self.blocks:
 
96
  x = self.final(x)
97
 
98
  if self.upsample_factor > 1:
99
+ if images is None:
100
+ raise ValueError(f'`images` cannot be `None` when the head\'s `upsample_factor > 1`!')
101
+ if patch_size is None:
102
+ raise ValueError(f'`patch_size` cannot be `None` when the head\'s `upsample_factor > 1`!')
103
+ h, w = tuple(d // patch_size for d in images.shape[-2:])
104
+ x = rearrange(x, 'b (h w) (u1 u2 c) -> b (h u1 w u2) c',
105
  h=h, w=w, u1=self.upsample_factor, u2=self.upsample_factor,
106
  c=self._real_output_dim)
107
 
 
123
  return state
124
 
125
 
126
+ def get_mlp_info_from_state(version: str, state: Dict[str, torch.Tensor], prefix: str = '', spectral_weights: bool = False):
127
  state = strip_prefix(state, prefix)
128
 
129
+ weight_suffix = 'weight' if not spectral_weights else 'parametrizations.weight.original'
130
+
131
  if version == 'v1':
132
+ hidden_dim, input_dim = state[f'fc1.{weight_suffix}'].shape
133
+ output_dim = state[f'fc2.{weight_suffix}'].shape[0]
134
 
135
  for num_inner in range(1000):
136
  k = f'inner.{num_inner}.0.weight'
137
  if k not in state:
138
  break
139
  elif version == 'v2':
140
+ hidden_dim, input_dim = state[f'fc1.{weight_suffix}'].shape
141
+ output_dim = state[f'final.2.{weight_suffix}'].shape[0]
142
 
143
  for num_inner in range(1000):
144
  k = f'blocks.{num_inner}.0.weight'
 
150
  return input_dim, hidden_dim, output_dim, num_inner
151
 
152
 
153
+ def create_mlp_from_config(version: str, input_dim: int, hidden_dim: int, output_dim: int, num_inner: int, **kwargs):
154
+ ret: nn.Module = MLP_FACTORY[version](input_dim, hidden_dim, output_dim, num_inner, from_config=True, **kwargs)
155
+
156
+ return ret
157
+
158
+
159
+ def create_mlp_from_state(version: str, state: Dict[str, torch.Tensor], prefix: str = '', spectral_weights: bool = False, **kwargs):
160
  state = strip_prefix(state, prefix)
161
 
162
+ input_dim, hidden_dim, output_dim, num_inner = get_mlp_info_from_state(version, state, spectral_weights=spectral_weights)
163
 
164
+ ret: nn.Module = create_mlp_from_config(version, input_dim, hidden_dim, output_dim, num_inner, **kwargs)
165
+
166
+ if spectral_weights:
167
+ enable_spectral_reparam(ret, init_norm_to_current=False, state_dict_guidance=state)
168
 
169
  ret.load_state_dict(state)
170
 
171
+ if spectral_weights:
172
+ disable_spectral_reparam(ret)
173
+
174
  return ret
cls_token.py CHANGED
@@ -5,6 +5,7 @@
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
@@ -14,7 +15,8 @@ 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
 
@@ -23,7 +25,9 @@ class ClsToken(nn.Module):
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
 
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 Optional
9
 
10
  import torch
11
  from torch import nn
 
15
  def __init__(self, ndim: int,
16
  num_tokens: int = 1,
17
  enabled: bool = True,
18
+ register_multiple: Optional[int] = None,
19
+ num_registers: Optional[int] = None,
20
  ):
21
  super().__init__()
22
 
 
25
  self.num_registers = 0
26
  self.num_tokens = num_tokens
27
  if enabled:
28
+ if num_registers:
29
+ self.num_registers = num_registers
30
+ elif register_multiple:
31
  self.num_registers = register_multiple - (num_tokens % register_multiple)
32
 
33
  scale = ndim ** -0.5
common.py CHANGED
@@ -38,6 +38,27 @@ RESOURCE_MAP = {
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",
@@ -66,6 +87,22 @@ RESOURCE_MAP = {
66
  max_resolution=2048,
67
  preferred_resolution=Resolution(512, 512),
68
  ),
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
69
  }
70
 
71
- DEFAULT_VERSION = "radio_v2.5-l"
 
38
  preferred_resolution=(768, 768),
39
  vitdet_num_global=4,
40
  ),
41
+ "radio_v2.5-h": RadioResource(
42
+ "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.5-h.pth.tar?download=true",
43
+ patch_size=16,
44
+ max_resolution=2048,
45
+ preferred_resolution=(768, 768),
46
+ vitdet_num_global=4,
47
+ ),
48
+ "radio_v2.5-h-norm": RadioResource(
49
+ "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.5-h-norm.pth.tar?download=true",
50
+ patch_size=16,
51
+ max_resolution=2048,
52
+ preferred_resolution=(768, 768),
53
+ vitdet_num_global=4,
54
+ ),
55
+ "radio_v2.5-g": RadioResource(
56
+ "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.5-g.pth.tar?download=true",
57
+ patch_size=14,
58
+ max_resolution=1792,
59
+ preferred_resolution=(896, 896),
60
+ vitdet_num_global=8,
61
+ ),
62
  # RADIO
63
  "radio_v2.1": RadioResource(
64
  "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.1_bf16.pth.tar?download=true",
 
87
  max_resolution=2048,
88
  preferred_resolution=Resolution(512, 512),
89
  ),
90
+ # C-RADIO
91
+ "c-radio_v2.5-g": RadioResource(
92
+ "https://huggingface.co/nvidia/C-RADIOv2-g/resolve/main/c-radio_v2-g_half.pth.tar",
93
+ patch_size=16,
94
+ max_resolution=2048,
95
+ preferred_resolution=(768, 768),
96
+ vitdet_num_global=8,
97
+ ),
98
+ "c-radio_v3-l": RadioResource(
99
+ # NOTE: Currently, this model cannot be loaded via TorchHub. Instead, use the transformers API at https://huggingface.co/nvidia/C-RADIOv3-L
100
+ # and accept the license terms.
101
+ "https://huggingface.co/nvidia/C-RADIOv3-L/resolve/main/c-radio-v3_l_half.pth.tar?download=true",
102
+ patch_size=16,
103
+ max_resolution=2048,
104
+ preferred_resolution=Resolution(512, 512),
105
+ ),
106
  }
107
 
108
+ DEFAULT_VERSION = "radio_v2.5-h"
config.json CHANGED
@@ -1,4 +1,5 @@
1
  {
 
2
  "adaptor_names": null,
3
  "architectures": [
4
  "RADIOModel"
@@ -39,7 +40,6 @@
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,
@@ -208,6 +208,10 @@
208
  "AutoConfig": "hf_model.RADIOConfig",
209
  "AutoModel": "hf_model.RADIOModel"
210
  },
 
 
 
 
211
  "max_resolution": 2048,
212
  "patch_size": 16,
213
  "preferred_resolution": [
@@ -215,7 +219,7 @@
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
  }
 
1
  {
2
+ "adaptor_configs": {},
3
  "adaptor_names": null,
4
  "architectures": [
5
  "RADIOModel"
 
40
  270
41
  ],
42
  "decay_rate": 0.1,
 
43
  "dist_bn": "reduce",
44
  "dist_norm_weight": 0.0,
45
  "distributed": true,
 
208
  "AutoConfig": "hf_model.RADIOConfig",
209
  "AutoModel": "hf_model.RADIOModel"
210
  },
211
+ "feature_normalizer_config": {
212
+ "embed_dim": 768
213
+ },
214
+ "inter_feature_normalizer_config": null,
215
  "max_resolution": 2048,
216
  "patch_size": 16,
217
  "preferred_resolution": [
 
219
  768
220
  ],
221
  "torch_dtype": "float32",
222
+ "transformers_version": "4.47.0.dev0",
223
  "version": "radio_v2.5-b",
224
  "vitdet_window_size": null
225
  }
dinov2_arch.py ADDED
@@ -0,0 +1,1016 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ #
3
+ # This source code is licensed under the Apache License, Version 2.0
4
+ # found in the LICENSE file in the root directory of this source tree.
5
+
6
+ # References:
7
+ # https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
8
+ # https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
9
+
10
+ # Nvidia
11
+ # NOTE: We re-define this model architecture primarily so that we don't have to worry about version compatibility breaking,
12
+ # but also because Huggingface does a string replace of `gamma` to something else when loading the model state,
13
+ # and this breaks loading of this model.
14
+
15
+ from enum import Enum
16
+ from functools import partial
17
+ import logging
18
+ import math
19
+ import os
20
+ import sys
21
+ from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
22
+ import warnings
23
+
24
+ import torch
25
+ from torch import nn
26
+ from torch.nn import functional as F
27
+ from torch.nn.init import trunc_normal_
28
+
29
+ _torch_has_sdpa = hasattr(F, 'scaled_dot_product_attention')
30
+
31
+
32
+ XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
33
+ try:
34
+ if XFORMERS_ENABLED:
35
+ from xformers.ops import fmha, scaled_index_add, index_select_cat, SwiGLU, memory_efficient_attention, unbind
36
+
37
+ XFORMERS_AVAILABLE = True
38
+ else:
39
+ raise ImportError
40
+ except ImportError:
41
+ XFORMERS_AVAILABLE = False
42
+
43
+
44
+ def make_2tuple(x):
45
+ if isinstance(x, tuple):
46
+ assert len(x) == 2
47
+ return x
48
+
49
+ assert isinstance(x, int)
50
+ return (x, x)
51
+
52
+
53
+ class PatchEmbed(nn.Module):
54
+ """
55
+ 2D image to patch embedding: (B,C,H,W) -> (B,N,D)
56
+
57
+ Args:
58
+ img_size: Image size.
59
+ patch_size: Patch token size.
60
+ in_chans: Number of input image channels.
61
+ embed_dim: Number of linear projection output channels.
62
+ norm_layer: Normalization layer.
63
+ """
64
+
65
+ def __init__(
66
+ self,
67
+ img_size: Union[int, Tuple[int, int]] = 224,
68
+ patch_size: Union[int, Tuple[int, int]] = 16,
69
+ in_chans: int = 3,
70
+ embed_dim: int = 768,
71
+ norm_layer: Optional[Callable] = None,
72
+ flatten_embedding: bool = True,
73
+ ) -> None:
74
+ super().__init__()
75
+
76
+ image_HW = make_2tuple(img_size)
77
+ patch_HW = make_2tuple(patch_size)
78
+ patch_grid_size = (
79
+ image_HW[0] // patch_HW[0],
80
+ image_HW[1] // patch_HW[1],
81
+ )
82
+
83
+ self.img_size = image_HW
84
+ self.patch_size = patch_HW
85
+ self.patches_resolution = patch_grid_size
86
+ self.num_patches = patch_grid_size[0] * patch_grid_size[1]
87
+
88
+ self.in_chans = in_chans
89
+ self.embed_dim = embed_dim
90
+
91
+ self.flatten_embedding = flatten_embedding
92
+
93
+ self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW)
94
+ self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
95
+
96
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
97
+ _, _, H, W = x.shape
98
+ patch_H, patch_W = self.patch_size
99
+
100
+ assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}"
101
+ assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}"
102
+
103
+ x = self.proj(x) # B C H W
104
+ H, W = x.size(2), x.size(3)
105
+ x = x.flatten(2).transpose(1, 2) # B HW C
106
+ x = self.norm(x)
107
+ if not self.flatten_embedding:
108
+ x = x.reshape(-1, H, W, self.embed_dim) # B H W C
109
+ return x
110
+
111
+ def flops(self) -> float:
112
+ Ho, Wo = self.patches_resolution
113
+ flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
114
+ if self.norm is not None:
115
+ flops += Ho * Wo * self.embed_dim
116
+ return flops
117
+
118
+
119
+ class Attention(nn.Module):
120
+ def __init__(
121
+ self,
122
+ dim: int,
123
+ num_heads: int = 8,
124
+ qkv_bias: bool = False,
125
+ proj_bias: bool = True,
126
+ attn_drop: float = 0.0,
127
+ proj_drop: float = 0.0,
128
+ ) -> None:
129
+ super().__init__()
130
+ self.num_heads = num_heads
131
+ head_dim = dim // num_heads
132
+ self.scale = head_dim**-0.5
133
+
134
+ self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
135
+ self.attn_drop = nn.Dropout(attn_drop)
136
+ self.proj = nn.Linear(dim, dim, bias=proj_bias)
137
+ self.proj_drop = nn.Dropout(proj_drop)
138
+
139
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
140
+ B, N, C = x.shape
141
+ qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
142
+
143
+ q, k, v = qkv[0], qkv[1], qkv[2]
144
+ if _torch_has_sdpa:
145
+ x = F.scaled_dot_product_attention(
146
+ q, k, v,
147
+ is_causal=False,
148
+ dropout_p=self.attn_drop.p if self.training else 0.,
149
+ scale=self.scale,
150
+ )
151
+ else:
152
+ q = q * self.scale
153
+ attn = q @ k.transpose(-2, -1)
154
+
155
+ attn = attn.softmax(dim=-1)
156
+ attn = self.attn_drop(attn)
157
+ x = attn @ v
158
+
159
+ x = x.transpose(1, 2).reshape(B, N, C)
160
+ x = self.proj(x)
161
+ x = self.proj_drop(x)
162
+ return x
163
+
164
+
165
+ class MemEffAttention(Attention):
166
+ def forward(self, x: torch.Tensor, attn_bias=None) -> torch.Tensor:
167
+ if not XFORMERS_AVAILABLE:
168
+ if attn_bias is not None:
169
+ raise AssertionError("xFormers is required for using nested tensors")
170
+ return super().forward(x)
171
+
172
+ B, N, C = x.shape
173
+ qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
174
+
175
+ q, k, v = unbind(qkv, 2)
176
+
177
+ x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
178
+ x = x.reshape([B, N, C])
179
+
180
+ x = self.proj(x)
181
+ x = self.proj_drop(x)
182
+ return x
183
+
184
+
185
+ class Mlp(nn.Module):
186
+ def __init__(
187
+ self,
188
+ in_features: int,
189
+ hidden_features: Optional[int] = None,
190
+ out_features: Optional[int] = None,
191
+ act_layer: Callable[..., nn.Module] = nn.GELU,
192
+ drop: float = 0.0,
193
+ bias: bool = True,
194
+ ) -> None:
195
+ super().__init__()
196
+ out_features = out_features or in_features
197
+ hidden_features = hidden_features or in_features
198
+ self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
199
+ self.act = act_layer()
200
+ self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
201
+ self.drop = nn.Dropout(drop)
202
+
203
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
204
+ x = self.fc1(x)
205
+ x = self.act(x)
206
+ x = self.drop(x)
207
+ x = self.fc2(x)
208
+ x = self.drop(x)
209
+ return x
210
+
211
+
212
+ class SwiGLUFFN(nn.Module):
213
+ def __init__(
214
+ self,
215
+ in_features: int,
216
+ hidden_features: Optional[int] = None,
217
+ out_features: Optional[int] = None,
218
+ act_layer: Callable[..., nn.Module] = None,
219
+ drop: float = 0.0,
220
+ bias: bool = True,
221
+ ) -> None:
222
+ super().__init__()
223
+ out_features = out_features or in_features
224
+ hidden_features = hidden_features or in_features
225
+ self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
226
+ self.w3 = nn.Linear(hidden_features, out_features, bias=bias)
227
+
228
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
229
+ x12 = self.w12(x)
230
+ x1, x2 = x12.chunk(2, dim=-1)
231
+ hidden = F.silu(x1) * x2
232
+ return self.w3(hidden)
233
+
234
+
235
+ if not XFORMERS_AVAILABLE:
236
+ SwiGLU = SwiGLUFFN
237
+
238
+
239
+ class SwiGLUFFNFused(SwiGLU):
240
+ def __init__(
241
+ self,
242
+ in_features: int,
243
+ hidden_features: Optional[int] = None,
244
+ out_features: Optional[int] = None,
245
+ act_layer: Callable[..., nn.Module] = None,
246
+ drop: float = 0.0,
247
+ bias: bool = True,
248
+ ) -> None:
249
+ out_features = out_features or in_features
250
+ hidden_features = hidden_features or in_features
251
+ hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
252
+ super().__init__(
253
+ in_features=in_features,
254
+ hidden_features=hidden_features,
255
+ out_features=out_features,
256
+ bias=bias,
257
+ )
258
+
259
+
260
+ def drop_path(x, drop_prob: float = 0.0, training: bool = False):
261
+ if drop_prob == 0.0 or not training:
262
+ return x
263
+ keep_prob = 1 - drop_prob
264
+ shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
265
+ random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
266
+ if keep_prob > 0.0:
267
+ random_tensor.div_(keep_prob)
268
+ output = x * random_tensor
269
+ return output
270
+
271
+
272
+ class DropPath(nn.Module):
273
+ """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
274
+
275
+ def __init__(self, drop_prob=None):
276
+ super(DropPath, self).__init__()
277
+ self.drop_prob = drop_prob
278
+
279
+ def forward(self, x):
280
+ return drop_path(x, self.drop_prob, self.training)
281
+
282
+
283
+ class LayerScale(nn.Module):
284
+ def __init__(
285
+ self,
286
+ dim: int,
287
+ init_values: Union[float, torch.Tensor] = 1e-5,
288
+ inplace: bool = False,
289
+ ) -> None:
290
+ super().__init__()
291
+ self.inplace = inplace
292
+ self.grandma = nn.Parameter(init_values * torch.ones(dim))
293
+
294
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
295
+ return x.mul_(self.grandma) if self.inplace else x * self.grandma
296
+
297
+ def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
298
+ # Huggingface is absurd and it will rename strings that contain `gamma`, which means that the normal DINO implementation
299
+ # of LayerScale won't work with HFHub. So we rename the variable to 'grandma', and support loading checkpoints in either
300
+ # format
301
+ key_a = f'{prefix}gamma'
302
+ key_b = f'{prefix}grandma'
303
+ if key_a in state_dict:
304
+ gamma = state_dict[key_a]
305
+ elif key_b in state_dict:
306
+ gamma = state_dict[key_b]
307
+ else:
308
+ if strict:
309
+ raise KeyError(f"Couldn't find the key {key_a} nor {key_b} in the state dict!")
310
+ else:
311
+ missing_keys.append(key_a)
312
+ missing_keys.append(key_b)
313
+ unexpected_keys.extend(state_dict.keys())
314
+ gamma = None
315
+
316
+ if gamma is not None:
317
+ self.grandma.data.copy_(gamma)
318
+
319
+ # return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
320
+
321
+
322
+ class Block(nn.Module):
323
+ def __init__(
324
+ self,
325
+ dim: int,
326
+ num_heads: int,
327
+ mlp_ratio: float = 4.0,
328
+ qkv_bias: bool = False,
329
+ proj_bias: bool = True,
330
+ ffn_bias: bool = True,
331
+ drop: float = 0.0,
332
+ attn_drop: float = 0.0,
333
+ init_values=None,
334
+ drop_path: float = 0.0,
335
+ act_layer: Callable[..., nn.Module] = nn.GELU,
336
+ norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
337
+ attn_class: Callable[..., nn.Module] = Attention,
338
+ ffn_layer: Callable[..., nn.Module] = Mlp,
339
+ ) -> None:
340
+ super().__init__()
341
+ # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}")
342
+ self.norm1 = norm_layer(dim)
343
+ self.attn = attn_class(
344
+ dim,
345
+ num_heads=num_heads,
346
+ qkv_bias=qkv_bias,
347
+ proj_bias=proj_bias,
348
+ attn_drop=attn_drop,
349
+ proj_drop=drop,
350
+ )
351
+ self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
352
+ self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
353
+
354
+ self.norm2 = norm_layer(dim)
355
+ mlp_hidden_dim = int(dim * mlp_ratio)
356
+ self.mlp = ffn_layer(
357
+ in_features=dim,
358
+ hidden_features=mlp_hidden_dim,
359
+ act_layer=act_layer,
360
+ drop=drop,
361
+ bias=ffn_bias,
362
+ )
363
+ self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
364
+ self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
365
+
366
+ self.sample_drop_ratio = drop_path
367
+
368
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
369
+ def attn_residual_func(x: torch.Tensor) -> torch.Tensor:
370
+ return self.ls1(self.attn(self.norm1(x)))
371
+
372
+ def ffn_residual_func(x: torch.Tensor) -> torch.Tensor:
373
+ return self.ls2(self.mlp(self.norm2(x)))
374
+
375
+ if self.training and self.sample_drop_ratio > 0.1:
376
+ # the overhead is compensated only for a drop path rate larger than 0.1
377
+ x = drop_add_residual_stochastic_depth(
378
+ x,
379
+ residual_func=attn_residual_func,
380
+ sample_drop_ratio=self.sample_drop_ratio,
381
+ )
382
+ x = drop_add_residual_stochastic_depth(
383
+ x,
384
+ residual_func=ffn_residual_func,
385
+ sample_drop_ratio=self.sample_drop_ratio,
386
+ )
387
+ elif self.training and self.sample_drop_ratio > 0.0:
388
+ x = x + self.drop_path1(attn_residual_func(x))
389
+ x = x + self.drop_path1(ffn_residual_func(x)) # FIXME: drop_path2
390
+ else:
391
+ x = x + attn_residual_func(x)
392
+ x = x + ffn_residual_func(x)
393
+ return x
394
+
395
+
396
+ class NestedTensorBlock(Block):
397
+ def forward_nested(self, x_list: List[torch.Tensor]) -> List[torch.Tensor]:
398
+ """
399
+ x_list contains a list of tensors to nest together and run
400
+ """
401
+ assert isinstance(self.attn, MemEffAttention)
402
+
403
+ if self.training and self.sample_drop_ratio > 0.0:
404
+
405
+ def attn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
406
+ return self.attn(self.norm1(x), attn_bias=attn_bias)
407
+
408
+ def ffn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
409
+ return self.mlp(self.norm2(x))
410
+
411
+ x_list = drop_add_residual_stochastic_depth_list(
412
+ x_list,
413
+ residual_func=attn_residual_func,
414
+ sample_drop_ratio=self.sample_drop_ratio,
415
+ scaling_vector=self.ls1.grandma if isinstance(self.ls1, LayerScale) else None,
416
+ )
417
+ x_list = drop_add_residual_stochastic_depth_list(
418
+ x_list,
419
+ residual_func=ffn_residual_func,
420
+ sample_drop_ratio=self.sample_drop_ratio,
421
+ scaling_vector=self.ls2.grandma if isinstance(self.ls1, LayerScale) else None,
422
+ )
423
+ return x_list
424
+ else:
425
+
426
+ def attn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
427
+ return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))
428
+
429
+ def ffn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
430
+ return self.ls2(self.mlp(self.norm2(x)))
431
+
432
+ attn_bias, x = get_attn_bias_and_cat(x_list)
433
+ x = x + attn_residual_func(x, attn_bias=attn_bias)
434
+ x = x + ffn_residual_func(x)
435
+ return attn_bias.split(x)
436
+
437
+ def forward(self, x_or_x_list):
438
+ if isinstance(x_or_x_list, torch.Tensor):
439
+ return super().forward(x_or_x_list)
440
+ elif isinstance(x_or_x_list, list):
441
+ if not XFORMERS_AVAILABLE:
442
+ raise AssertionError("xFormers is required for using nested tensors")
443
+ return self.forward_nested(x_or_x_list)
444
+ else:
445
+ raise AssertionError
446
+
447
+
448
+ def drop_add_residual_stochastic_depth(
449
+ x: torch.Tensor,
450
+ residual_func: Callable[[torch.Tensor], torch.Tensor],
451
+ sample_drop_ratio: float = 0.0,
452
+ ) -> torch.Tensor:
453
+ # 1) extract subset using permutation
454
+ b, n, d = x.shape
455
+ sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
456
+ brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
457
+ x_subset = x[brange]
458
+
459
+ # 2) apply residual_func to get residual
460
+ residual = residual_func(x_subset)
461
+
462
+ x_flat = x.flatten(1)
463
+ residual = residual.flatten(1)
464
+
465
+ residual_scale_factor = b / sample_subset_size
466
+
467
+ # 3) add the residual
468
+ x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
469
+ return x_plus_residual.view_as(x)
470
+
471
+
472
+ def get_branges_scales(x, sample_drop_ratio=0.0):
473
+ b, n, d = x.shape
474
+ sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
475
+ brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
476
+ residual_scale_factor = b / sample_subset_size
477
+ return brange, residual_scale_factor
478
+
479
+
480
+ def add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None):
481
+ if scaling_vector is None:
482
+ x_flat = x.flatten(1)
483
+ residual = residual.flatten(1)
484
+ x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
485
+ else:
486
+ x_plus_residual = scaled_index_add(
487
+ x, brange, residual.to(dtype=x.dtype), scaling=scaling_vector, alpha=residual_scale_factor
488
+ )
489
+ return x_plus_residual
490
+
491
+
492
+ attn_bias_cache: Dict[Tuple, Any] = {}
493
+
494
+
495
+ def get_attn_bias_and_cat(x_list, branges=None):
496
+ """
497
+ this will perform the index select, cat the tensors, and provide the attn_bias from cache
498
+ """
499
+ batch_sizes = [b.shape[0] for b in branges] if branges is not None else [x.shape[0] for x in x_list]
500
+ all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))
501
+ if all_shapes not in attn_bias_cache.keys():
502
+ seqlens = []
503
+ for b, x in zip(batch_sizes, x_list):
504
+ for _ in range(b):
505
+ seqlens.append(x.shape[1])
506
+ attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)
507
+ attn_bias._batch_sizes = batch_sizes
508
+ attn_bias_cache[all_shapes] = attn_bias
509
+
510
+ if branges is not None:
511
+ cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(1, -1, x_list[0].shape[-1])
512
+ else:
513
+ tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)
514
+ cat_tensors = torch.cat(tensors_bs1, dim=1)
515
+
516
+ return attn_bias_cache[all_shapes], cat_tensors
517
+
518
+
519
+ def drop_add_residual_stochastic_depth_list(
520
+ x_list: List[torch.Tensor],
521
+ residual_func: Callable[[torch.Tensor, Any], torch.Tensor],
522
+ sample_drop_ratio: float = 0.0,
523
+ scaling_vector=None,
524
+ ) -> torch.Tensor:
525
+ # 1) generate random set of indices for dropping samples in the batch
526
+ branges_scales = [get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list]
527
+ branges = [s[0] for s in branges_scales]
528
+ residual_scale_factors = [s[1] for s in branges_scales]
529
+
530
+ # 2) get attention bias and index+concat the tensors
531
+ attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)
532
+
533
+ # 3) apply residual_func to get residual, and split the result
534
+ residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias)) # type: ignore
535
+
536
+ outputs = []
537
+ for x, brange, residual, residual_scale_factor in zip(x_list, branges, residual_list, residual_scale_factors):
538
+ outputs.append(add_residual(x, brange, residual, residual_scale_factor, scaling_vector).view_as(x))
539
+ return outputs
540
+
541
+
542
+ def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module:
543
+ if not depth_first and include_root:
544
+ fn(module=module, name=name)
545
+ for child_name, child_module in module.named_children():
546
+ child_name = ".".join((name, child_name)) if name else child_name
547
+ named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True)
548
+ if depth_first and include_root:
549
+ fn(module=module, name=name)
550
+ return module
551
+
552
+
553
+ class BlockChunk(nn.ModuleList):
554
+ def forward(self, x):
555
+ for b in self:
556
+ x = b(x)
557
+ return x
558
+
559
+
560
+ class DinoVisionTransformer(nn.Module):
561
+ def __init__(
562
+ self,
563
+ img_size=224,
564
+ patch_size=16,
565
+ in_chans=3,
566
+ embed_dim=768,
567
+ depth=12,
568
+ num_heads=12,
569
+ mlp_ratio=4.0,
570
+ qkv_bias=True,
571
+ ffn_bias=True,
572
+ proj_bias=True,
573
+ drop_path_rate=0.0,
574
+ drop_path_uniform=False,
575
+ init_values=None, # for layerscale: None or 0 => no layerscale
576
+ embed_layer=PatchEmbed,
577
+ act_layer=nn.GELU,
578
+ block_fn=Block,
579
+ ffn_layer="mlp",
580
+ block_chunks=1,
581
+ num_register_tokens=0,
582
+ interpolate_antialias=False,
583
+ interpolate_offset=0.1,
584
+ ):
585
+ """
586
+ Args:
587
+ img_size (int, tuple): input image size
588
+ patch_size (int, tuple): patch size
589
+ in_chans (int): number of input channels
590
+ embed_dim (int): embedding dimension
591
+ depth (int): depth of transformer
592
+ num_heads (int): number of attention heads
593
+ mlp_ratio (int): ratio of mlp hidden dim to embedding dim
594
+ qkv_bias (bool): enable bias for qkv if True
595
+ proj_bias (bool): enable bias for proj in attn if True
596
+ ffn_bias (bool): enable bias for ffn if True
597
+ drop_path_rate (float): stochastic depth rate
598
+ drop_path_uniform (bool): apply uniform drop rate across blocks
599
+ weight_init (str): weight init scheme
600
+ init_values (float): layer-scale init values
601
+ embed_layer (nn.Module): patch embedding layer
602
+ act_layer (nn.Module): MLP activation layer
603
+ block_fn (nn.Module): transformer block class
604
+ ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity"
605
+ block_chunks: (int) split block sequence into block_chunks units for FSDP wrap
606
+ num_register_tokens: (int) number of extra cls tokens (so-called "registers")
607
+ interpolate_antialias: (str) flag to apply anti-aliasing when interpolating positional embeddings
608
+ interpolate_offset: (float) work-around offset to apply when interpolating positional embeddings
609
+ """
610
+ super().__init__()
611
+ norm_layer = partial(nn.LayerNorm, eps=1e-6)
612
+
613
+ self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models
614
+ self.num_tokens = 1
615
+ self.n_blocks = depth
616
+ self.num_heads = num_heads
617
+ self.patch_size = patch_size
618
+ self.num_register_tokens = num_register_tokens
619
+ self.interpolate_antialias = interpolate_antialias
620
+ self.interpolate_offset = interpolate_offset
621
+
622
+ self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
623
+ num_patches = self.patch_embed.num_patches
624
+
625
+ self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
626
+ self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
627
+ assert num_register_tokens >= 0
628
+ self.register_tokens = (
629
+ nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None
630
+ )
631
+
632
+ if drop_path_uniform is True:
633
+ dpr = [drop_path_rate] * depth
634
+ else:
635
+ dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule
636
+
637
+ if ffn_layer == "mlp":
638
+ ffn_layer = Mlp
639
+ elif ffn_layer == "swiglufused" or ffn_layer == "swiglu":
640
+ ffn_layer = SwiGLUFFNFused
641
+ elif ffn_layer == "identity":
642
+ def f(*args, **kwargs):
643
+ return nn.Identity()
644
+
645
+ ffn_layer = f
646
+ else:
647
+ raise NotImplementedError
648
+
649
+ blocks_list = [
650
+ block_fn(
651
+ dim=embed_dim,
652
+ num_heads=num_heads,
653
+ mlp_ratio=mlp_ratio,
654
+ qkv_bias=qkv_bias,
655
+ proj_bias=proj_bias,
656
+ ffn_bias=ffn_bias,
657
+ drop_path=dpr[i],
658
+ norm_layer=norm_layer,
659
+ act_layer=act_layer,
660
+ ffn_layer=ffn_layer,
661
+ init_values=init_values,
662
+ )
663
+ for i in range(depth)
664
+ ]
665
+ if block_chunks > 0:
666
+ self.chunked_blocks = True
667
+ chunked_blocks = []
668
+ chunksize = depth // block_chunks
669
+ for i in range(0, depth, chunksize):
670
+ # this is to keep the block index consistent if we chunk the block list
671
+ chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize])
672
+ self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])
673
+ else:
674
+ self.chunked_blocks = False
675
+ self.blocks = nn.ModuleList(blocks_list)
676
+
677
+ self.norm = norm_layer(embed_dim)
678
+ self.head = nn.Identity()
679
+
680
+ self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))
681
+
682
+ def interpolate_pos_encoding(self, x, w, h):
683
+ previous_dtype = x.dtype
684
+ npatch = x.shape[1] - 1
685
+ N = self.pos_embed.shape[1] - 1
686
+ if npatch == N and w == h:
687
+ return self.pos_embed
688
+ pos_embed = self.pos_embed.float()
689
+ class_pos_embed = pos_embed[:, 0]
690
+ patch_pos_embed = pos_embed[:, 1:]
691
+ dim = x.shape[-1]
692
+ w0 = w // self.patch_size
693
+ h0 = h // self.patch_size
694
+ M = int(math.sqrt(N)) # Recover the number of patches in each dimension
695
+ assert N == M * M
696
+ kwargs = {}
697
+ if self.interpolate_offset:
698
+ # Historical kludge: add a small number to avoid floating point error in the interpolation, see https://github.com/facebookresearch/dino/issues/8
699
+ # Note: still needed for backward-compatibility, the underlying operators are using both output size and scale factors
700
+ sx = float(w0 + self.interpolate_offset) / M
701
+ sy = float(h0 + self.interpolate_offset) / M
702
+ kwargs["scale_factor"] = (sx, sy)
703
+ else:
704
+ # Simply specify an output size instead of a scale factor
705
+ kwargs["size"] = (w0, h0)
706
+ patch_pos_embed = nn.functional.interpolate(
707
+ patch_pos_embed.reshape(1, M, M, dim).permute(0, 3, 1, 2),
708
+ mode="bicubic",
709
+ antialias=self.interpolate_antialias,
710
+ **kwargs,
711
+ )
712
+ assert (w0, h0) == patch_pos_embed.shape[-2:]
713
+ patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
714
+ return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)
715
+
716
+ def prepare_tokens_with_masks(self, x, masks=None):
717
+ B, nc, w, h = x.shape
718
+ x = self.patch_embed(x)
719
+ if masks is not None:
720
+ x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
721
+
722
+ x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
723
+ x = x + self.interpolate_pos_encoding(x, w, h)
724
+
725
+ if self.register_tokens is not None:
726
+ x = torch.cat(
727
+ (
728
+ x[:, :1],
729
+ self.register_tokens.expand(x.shape[0], -1, -1),
730
+ x[:, 1:],
731
+ ),
732
+ dim=1,
733
+ )
734
+
735
+ return x
736
+
737
+ def forward_features_list(self, x_list, masks_list):
738
+ x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]
739
+ for blk in self.blocks:
740
+ x = blk(x)
741
+
742
+ all_x = x
743
+ output = []
744
+ for x, masks in zip(all_x, masks_list):
745
+ x_norm = self.norm(x)
746
+ output.append(
747
+ {
748
+ "x_norm_clstoken": x_norm[:, 0],
749
+ "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
750
+ "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
751
+ "x_prenorm": x,
752
+ "masks": masks,
753
+ }
754
+ )
755
+ return output
756
+
757
+ def forward_features(self, x, masks=None):
758
+ if isinstance(x, list):
759
+ return self.forward_features_list(x, masks)
760
+
761
+ x = self.prepare_tokens_with_masks(x, masks)
762
+
763
+ for blk in self.blocks:
764
+ x = blk(x)
765
+
766
+ x_norm = self.norm(x)
767
+ return {
768
+ "x_norm_clstoken": x_norm[:, 0],
769
+ "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
770
+ "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
771
+ "x_prenorm": x,
772
+ "masks": masks,
773
+ }
774
+
775
+ def _get_intermediate_layers_not_chunked(self, x, n=1):
776
+ x = self.prepare_tokens_with_masks(x)
777
+ # If n is an int, take the n last blocks. If it's a list, take them
778
+ output, total_block_len = [], len(self.blocks)
779
+ blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
780
+ for i, blk in enumerate(self.blocks):
781
+ x = blk(x)
782
+ if i in blocks_to_take:
783
+ output.append(x)
784
+ assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
785
+ return output
786
+
787
+ def _get_intermediate_layers_chunked(self, x, n=1):
788
+ x = self.prepare_tokens_with_masks(x)
789
+ output, i, total_block_len = [], 0, len(self.blocks[-1])
790
+ # If n is an int, take the n last blocks. If it's a list, take them
791
+ blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
792
+ for block_chunk in self.blocks:
793
+ for blk in block_chunk[i:]: # Passing the nn.Identity()
794
+ x = blk(x)
795
+ if i in blocks_to_take:
796
+ output.append(x)
797
+ i += 1
798
+ assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
799
+ return output
800
+
801
+ def get_intermediate_layers(
802
+ self,
803
+ x: torch.Tensor,
804
+ n: Union[int, Sequence] = 1, # Layers or n last layers to take
805
+ reshape: bool = False,
806
+ return_class_token: bool = False,
807
+ norm=True,
808
+ ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:
809
+ if self.chunked_blocks:
810
+ outputs = self._get_intermediate_layers_chunked(x, n)
811
+ else:
812
+ outputs = self._get_intermediate_layers_not_chunked(x, n)
813
+ if norm:
814
+ outputs = [self.norm(out) for out in outputs]
815
+ class_tokens = [out[:, 0] for out in outputs]
816
+ outputs = [out[:, 1 + self.num_register_tokens :] for out in outputs]
817
+ if reshape:
818
+ B, _, w, h = x.shape
819
+ outputs = [
820
+ out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous()
821
+ for out in outputs
822
+ ]
823
+ if return_class_token:
824
+ return tuple(zip(outputs, class_tokens))
825
+ return tuple(outputs)
826
+
827
+ def forward(self, *args, is_training=False, **kwargs):
828
+ ret = self.forward_features(*args, **kwargs)
829
+ if is_training:
830
+ return ret
831
+ else:
832
+ return self.head(ret["x_norm_clstoken"])
833
+
834
+
835
+ def vit_small(patch_size=16, num_register_tokens=0, **kwargs):
836
+ model = DinoVisionTransformer(
837
+ patch_size=patch_size,
838
+ embed_dim=384,
839
+ depth=12,
840
+ num_heads=6,
841
+ mlp_ratio=4,
842
+ block_fn=partial(Block, attn_class=MemEffAttention),
843
+ num_register_tokens=num_register_tokens,
844
+ **kwargs,
845
+ )
846
+ return model
847
+
848
+
849
+ def vit_base(patch_size=16, num_register_tokens=0, **kwargs):
850
+ model = DinoVisionTransformer(
851
+ patch_size=patch_size,
852
+ embed_dim=768,
853
+ depth=12,
854
+ num_heads=12,
855
+ mlp_ratio=4,
856
+ block_fn=partial(Block, attn_class=MemEffAttention),
857
+ num_register_tokens=num_register_tokens,
858
+ **kwargs,
859
+ )
860
+ return model
861
+
862
+
863
+ def vit_large(patch_size=16, num_register_tokens=0, **kwargs):
864
+ model = DinoVisionTransformer(
865
+ patch_size=patch_size,
866
+ embed_dim=1024,
867
+ depth=24,
868
+ num_heads=16,
869
+ mlp_ratio=4,
870
+ block_fn=partial(Block, attn_class=MemEffAttention),
871
+ num_register_tokens=num_register_tokens,
872
+ **kwargs,
873
+ )
874
+ return model
875
+
876
+
877
+ def vit_giant2(patch_size=16, num_register_tokens=0, **kwargs):
878
+ """
879
+ Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64
880
+ """
881
+ model = DinoVisionTransformer(
882
+ patch_size=patch_size,
883
+ embed_dim=1536,
884
+ depth=40,
885
+ num_heads=24,
886
+ mlp_ratio=4,
887
+ block_fn=partial(Block, attn_class=MemEffAttention),
888
+ num_register_tokens=num_register_tokens,
889
+ **kwargs,
890
+ )
891
+ return model
892
+
893
+
894
+ class Weights(Enum):
895
+ LVD142M = "LVD142M"
896
+
897
+
898
+ def _make_dinov2_model(
899
+ *,
900
+ arch_name: str = "vit_large",
901
+ img_size: int = 518,
902
+ patch_size: int = 14,
903
+ init_values: float = 1.0,
904
+ ffn_layer: str = "mlp",
905
+ block_chunks: int = 0,
906
+ num_register_tokens: int = 0,
907
+ interpolate_antialias: bool = False,
908
+ interpolate_offset: float = 0.1,
909
+ weights: Union[Weights, str] = Weights.LVD142M,
910
+ **kwargs,
911
+ ):
912
+ if isinstance(weights, str):
913
+ try:
914
+ weights = Weights[weights]
915
+ except KeyError:
916
+ raise AssertionError(f"Unsupported weights: {weights}")
917
+
918
+ vit_kwargs = dict(
919
+ img_size=img_size,
920
+ patch_size=patch_size,
921
+ init_values=init_values,
922
+ ffn_layer=ffn_layer,
923
+ block_chunks=block_chunks,
924
+ num_register_tokens=num_register_tokens,
925
+ interpolate_antialias=interpolate_antialias,
926
+ interpolate_offset=interpolate_offset,
927
+ )
928
+ vit_kwargs.update(**kwargs)
929
+ model = sys.modules[__name__].__dict__[arch_name](**vit_kwargs)
930
+
931
+ return model
932
+
933
+
934
+ def dinov2_vits14(**kwargs):
935
+ """
936
+ DINOv2 ViT-S/14 model (optionally) pretrained on the LVD-142M dataset.
937
+ """
938
+ return _make_dinov2_model(arch_name="vit_small", **kwargs)
939
+
940
+
941
+ def dinov2_vitb14(**kwargs):
942
+ """
943
+ DINOv2 ViT-B/14 model (optionally) pretrained on the LVD-142M dataset.
944
+ """
945
+ return _make_dinov2_model(arch_name="vit_base", **kwargs)
946
+
947
+
948
+ def dinov2_vitl14(**kwargs):
949
+ """
950
+ DINOv2 ViT-L/14 model (optionally) pretrained on the LVD-142M dataset.
951
+ """
952
+ return _make_dinov2_model(arch_name="vit_large", **kwargs)
953
+
954
+
955
+ def dinov2_vitg14(**kwargs):
956
+ """
957
+ DINOv2 ViT-g/14 model (optionally) pretrained on the LVD-142M dataset.
958
+ """
959
+ return _make_dinov2_model(
960
+ arch_name="vit_giant2",
961
+ ffn_layer="swiglufused",
962
+ **kwargs,
963
+ )
964
+
965
+
966
+ def dinov2_vits14_reg(**kwargs):
967
+ """
968
+ DINOv2 ViT-S/14 model with registers (optionally) pretrained on the LVD-142M dataset.
969
+ """
970
+ return _make_dinov2_model(
971
+ arch_name="vit_small",
972
+ num_register_tokens=4,
973
+ interpolate_antialias=True,
974
+ interpolate_offset=0.0,
975
+ **kwargs,
976
+ )
977
+
978
+
979
+ def dinov2_vitb14_reg(**kwargs):
980
+ """
981
+ DINOv2 ViT-B/14 model with registers (optionally) pretrained on the LVD-142M dataset.
982
+ """
983
+ return _make_dinov2_model(
984
+ arch_name="vit_base",
985
+ num_register_tokens=4,
986
+ interpolate_antialias=True,
987
+ interpolate_offset=0.0,
988
+ **kwargs,
989
+ )
990
+
991
+
992
+ def dinov2_vitl14_reg(**kwargs):
993
+ """
994
+ DINOv2 ViT-L/14 model with registers (optionally) pretrained on the LVD-142M dataset.
995
+ """
996
+ return _make_dinov2_model(
997
+ arch_name="vit_large",
998
+ num_register_tokens=4,
999
+ interpolate_antialias=True,
1000
+ interpolate_offset=0.0,
1001
+ **kwargs,
1002
+ )
1003
+
1004
+
1005
+ def dinov2_vitg14_reg(**kwargs):
1006
+ """
1007
+ DINOv2 ViT-g/14 model with registers (optionally) pretrained on the LVD-142M dataset.
1008
+ """
1009
+ return _make_dinov2_model(
1010
+ arch_name="vit_giant2",
1011
+ ffn_layer="swiglufused",
1012
+ num_register_tokens=4,
1013
+ interpolate_antialias=True,
1014
+ interpolate_offset=0.0,
1015
+ **kwargs,
1016
+ )
dual_hybrid_vit.py ADDED
@@ -0,0 +1,213 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from logging import getLogger
2
+ from typing import Tuple
3
+
4
+ import torch
5
+ from torch import nn
6
+ from torch.nn import functional as F
7
+
8
+ from timm.models import register_model
9
+ from timm.models import vision_transformer as tvit
10
+ from timm.models import convnext as tconv
11
+
12
+ from einops import rearrange
13
+
14
+ from . import extra_timm_models as et
15
+
16
+
17
+ class Fuser(nn.Module):
18
+ def __init__(self, src_dim: int, tgt_dim: int, gated: bool = True):
19
+ super().__init__()
20
+ self.gated = gated
21
+
22
+ mid_dim = max(src_dim, tgt_dim) * 2
23
+
24
+ self.fwd = nn.Sequential(
25
+ nn.Conv2d(src_dim, mid_dim, kernel_size=3, stride=1, padding=1),
26
+ nn.GELU(),
27
+ nn.Conv2d(mid_dim, tgt_dim * (2 if gated else 1), kernel_size=3, stride=1, padding=1),
28
+ )
29
+
30
+ def forward(self, src: torch.Tensor, tgt: torch.Tensor) -> torch.Tensor:
31
+ if src.ndim == 3:
32
+ shape = tgt.shape[-2:]
33
+ else:
34
+ shape = src.shape[-2:]
35
+
36
+ nd = shape[0] * shape[1]
37
+
38
+ if src.ndim == 3:
39
+ src = src[:, -nd:].reshape(src.shape[0], src.shape[2], *shape)
40
+
41
+ if tgt.ndim == 3:
42
+ tgt_pre = tgt[:, :-nd]
43
+ tgt = tgt[:, -nd:].reshape(tgt.shape[0], tgt.shape[2], *shape)
44
+ else:
45
+ tgt_pre = None
46
+
47
+ pred = self.fwd(src)
48
+
49
+ if self.gated:
50
+ g, pred = torch.chunk(pred, 2, dim=1)
51
+
52
+ g = F.sigmoid(g)
53
+
54
+ pred = g * pred
55
+
56
+ tgt = tgt + pred
57
+
58
+ if tgt_pre is not None:
59
+ tgt = rearrange(tgt, 'b c h w -> b (h w) c')
60
+ tgt = torch.cat([tgt_pre, tgt], dim=1)
61
+
62
+ return tgt
63
+
64
+
65
+ class AttnDownsample(nn.Module):
66
+ def __init__(self, dim: int, window_size: int, num_heads: int = 16):
67
+ super().__init__()
68
+ self.q = nn.Parameter(torch.randn(1, num_heads, 1, dim // num_heads) * 0.01)
69
+ self.kv = nn.Linear(dim, dim * 2)
70
+ self.proj = nn.Linear(dim, dim)
71
+ self.window_size = window_size
72
+ self.num_heads = num_heads
73
+ self.head_dim = dim // num_heads
74
+ self.scale = self.head_dim ** -0.5
75
+
76
+ def forward(self, x: torch.Tensor, twod_shape: Tuple[int, int]) -> torch.Tensor:
77
+ ntok = twod_shape[0] * twod_shape[1]
78
+ x_pre = x[:, :-ntok]
79
+
80
+ B = x.shape[0]
81
+ ds_hw = tuple(s // self.window_size for s in twod_shape)
82
+
83
+ x_spat = rearrange(
84
+ x[:, -ntok:],
85
+ 'b (h d1 w d2) c -> (b h w) (d1 d2) c',
86
+ h=ds_hw[0], w=ds_hw[1],
87
+ d1=self.window_size, d2=self.window_size,
88
+ )
89
+
90
+ B, N, C = x_spat.shape
91
+
92
+ k, v = self.kv(x_spat).reshape(B, N, 2, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
93
+
94
+ q = (self.q * self.scale).expand(B, -1, -1, -1)
95
+ attn = q @ k.transpose(-2, -1)
96
+ attn = F.softmax(attn, dim=-1)
97
+ x = attn @ v
98
+
99
+ x = x.transpose(1, 2).reshape(B, C)
100
+ x = self.proj(x)
101
+
102
+ x = rearrange(x, '(b h w) c -> b (h w) c', b=x_pre.shape[0], h=ds_hw[0], w=ds_hw[1])
103
+
104
+ x = torch.cat([x_pre, x], dim=1)
105
+ return x
106
+
107
+
108
+ class HybridModel(nn.Module):
109
+ def __init__(self, vit: tvit.VisionTransformer, conv: tconv.ConvNeXt, pretrained: bool = False,
110
+ concatenate: bool = False, **kwargs):
111
+ super().__init__()
112
+ self.conv = conv
113
+ self.vit = vit
114
+ self.concatenate = concatenate
115
+
116
+ conv.stages = nn.ModuleList(conv.stages)
117
+ vit.blocks = nn.ModuleList(vit.blocks)
118
+
119
+ self._half_vit_idx = len(vit.blocks) // 2 + 1
120
+
121
+ self._half_conv_idx = None
122
+ x = torch.empty(1, 3, 256, 256)
123
+ x = self.conv.stem(x)
124
+ for i in range(len(conv.stages)):
125
+ x = conv.stages[i](x)
126
+ if self._half_conv_idx is None and x.shape[-2:] == (16, 16):
127
+ self._half_conv_idx = i + 1
128
+ half_conv_dim = x.shape[1]
129
+ final_conv_dim = x.shape[1]
130
+
131
+ self.vit_to_conv_fusion = Fuser(vit.embed_dim, half_conv_dim)
132
+ self.conv_to_vit_fusion = Fuser(half_conv_dim, vit.embed_dim)
133
+ self.vit_ds = AttnDownsample(vit.embed_dim, window_size=2)
134
+
135
+ embed_dim = vit.embed_dim + (final_conv_dim if concatenate else 0)
136
+ if not concatenate:
137
+ self.final_fuse = Fuser(final_conv_dim, vit.embed_dim, gated=False)
138
+ self.final_block = tvit.Block(embed_dim, num_heads=16)
139
+
140
+ self.embed_dim = embed_dim
141
+
142
+ @property
143
+ def patch_size(self):
144
+ return 32
145
+
146
+ @property
147
+ def no_fsdp_wrap_types(self):
148
+ return {tvit.VisionTransformer, tconv.ConvNeXt}
149
+
150
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
151
+ return self.forward_features(x)
152
+
153
+ def forward_features(self, x: torch.Tensor) -> torch.Tensor:
154
+ y_vit = self.vit.patch_generator(x)
155
+
156
+ for i in range(self._half_vit_idx):
157
+ y_vit = self.vit.blocks[i](y_vit)
158
+
159
+ y_conv = self.conv.stem(x)
160
+ for i in range(self._half_conv_idx):
161
+ y_conv = self.conv.stages[i](y_conv)
162
+
163
+ y_vit, y_conv = self.conv_to_vit_fusion(y_conv, y_vit), self.vit_to_conv_fusion(y_vit, y_conv)
164
+
165
+ y_vit = self.vit_ds(y_vit, y_conv.shape[-2:])
166
+
167
+ for i in range(self._half_vit_idx, len(self.vit.blocks)):
168
+ y_vit = self.vit.blocks[i](y_vit)
169
+
170
+ for i in range(self._half_conv_idx, len(self.conv.stages)):
171
+ y_conv = self.conv.stages[i](y_conv)
172
+
173
+ if self.concatenate:
174
+ y_conv = rearrange(y_conv, 'b c h w -> b (h w) c')
175
+ # Average pool across the board, and replicate for each cls/register token
176
+ conv_summary = y_conv.mean(dim=1, keepdim=True).expand(-1, self.vit.patch_generator.num_cls_patches, -1)
177
+ y_conv = torch.cat([conv_summary, y_conv], dim=1)
178
+ y = torch.cat([y_vit, y_conv], dim=2)
179
+ else:
180
+ y = self.final_fuse(y_conv, y_vit)
181
+ y = self.final_block(y)
182
+
183
+ summary = y[:, :self.vit.patch_generator.num_cls_tokens]
184
+ features = y[:, self.vit.patch_generator.num_cls_patches:]
185
+
186
+ return summary, features
187
+
188
+
189
+ @register_model
190
+ def hybrid_base(pretrained=False, concatenate: bool = False, weight_init: str = 'skip', **kwargs):
191
+ cfg = dict(num_classes=0, **kwargs)
192
+ conv = tconv.convnextv2_base(pretrained=pretrained, **cfg)
193
+ vit = tvit.vit_base_patch16_224(pretrained=pretrained, weight_init=weight_init, **cfg)
194
+
195
+ return HybridModel(vit, conv, pretrained, concatenate=concatenate)
196
+
197
+
198
+ @register_model
199
+ def hybrid_large(pretrained=False, concatenate: bool = False, weight_init: str = 'skip', **kwargs):
200
+ cfg = dict(num_classes=0, **kwargs)
201
+ conv = tconv.convnextv2_large(pretrained=pretrained, **cfg)
202
+ vit = tvit.vit_large_patch16_224(pretrained=pretrained, weight_init=weight_init, **cfg)
203
+
204
+ return HybridModel(vit, conv, pretrained, concatenate=concatenate)
205
+
206
+
207
+ @register_model
208
+ def hybrid_huge(pretrained=False, concatenate: bool = False, weight_init: str = 'skip', **kwargs):
209
+ cfg = dict(num_classes=0, **kwargs)
210
+ conv = tconv.convnextv2_huge(pretrained=pretrained, **cfg)
211
+ vit = et.vit_huge_patch16_224(pretrained=pretrained, weight_init=weight_init, **cfg)
212
+
213
+ return HybridModel(vit, conv, pretrained, concatenate=concatenate)
enable_cpe_support.py CHANGED
@@ -14,12 +14,17 @@ 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)
@@ -42,86 +47,36 @@ def _take_indices(
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!")
@@ -144,6 +99,7 @@ def enable_cpe(model: nn.Module,
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
@@ -151,8 +107,64 @@ def enable_cpe(model: nn.Module,
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)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
14
 
15
  from timm.models import VisionTransformer, checkpoint_seq
16
 
17
+ from .feature_normalizer import IntermediateFeatureNormalizerBase, NullIntermediateFeatureNormalizer
18
+
19
+ from .extra_models import DinoWrapper
20
  from .vit_patch_generator import ViTPatchGenerator
21
+ from .forward_intermediates import forward_intermediates
22
+ from .dual_hybrid_vit import HybridModel
23
 
24
 
25
  def _forward_cpe(self: VisionTransformer, x: torch.Tensor) -> torch.Tensor:
26
  x = self.patch_generator(x)
27
+ if getattr(self, 'grad_checkpointing', False) and not torch.jit.is_scripting():
28
  x = checkpoint_seq(self.blocks, x)
29
  else:
30
  x = self.blocks(x)
 
47
  def _forward_intermediates_cpe(
48
  self,
49
  x: torch.Tensor,
 
 
50
  norm: bool = False,
51
+ **kwargs,
 
 
 
52
  ) -> Union[List[torch.Tensor], Tuple[torch.Tensor, List[torch.Tensor]]]:
53
+ return forward_intermediates(
54
+ self,
55
+ patch_extractor=self.patch_generator,
56
+ num_summary_tokens=self.patch_generator.num_skip,
57
+ num_cls_tokens=self.patch_generator.num_cls_tokens,
58
+ norm=self.norm if norm else lambda y: y,
59
+ x=x,
60
+ **kwargs,
61
+ )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
62
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
63
 
64
+ def _forward_cpe_dinov2(self: DinoWrapper, x: torch.Tensor) -> torch.Tensor:
65
+ y = _forward_cpe(self.inner, x)
66
+
67
+ return y[:, 0], y[:, self.num_summary_tokens:]
68
+
69
+
70
+ def _forward_intermediates_cpe_dinov2(self: DinoWrapper, *args, **kwargs):
71
+ return _forward_intermediates_cpe(self.inner, *args, **kwargs)
72
+
73
+
74
+ def _enable_cpe_for_timm_vit(model: VisionTransformer,
75
+ max_img_size: Union[int, Tuple[int, int]] = 1024,
76
+ num_cls_tokens: int = 1,
77
+ pos_dropout: float = 0.1,
78
+ register_multiple: int = Optional[None],
79
+ num_registers: int = Optional[None],
80
  ):
81
  if not isinstance(model, VisionTransformer):
82
  raise ValueError("CPE only support for VisionTransformer models!")
 
99
  pos_dropout=pos_dropout,
100
  num_cls_tokens=num_cls_tokens,
101
  register_multiple=register_multiple,
102
+ num_registers=num_registers,
103
  )
104
 
105
  model.patch_generator = patch_generator
 
107
  model.cls_token = None
108
  model.pos_embed = None
109
  model.pos_drop = None
110
+ model.patch_size = patch_size
111
  model.num_cls_tokens = num_cls_tokens
112
  model.num_registers = patch_generator.num_registers
113
 
114
  model.forward_features = MethodType(_forward_cpe, model)
115
  model.forward_intermediates = MethodType(_forward_intermediates_cpe, model)
116
+
117
+
118
+ def _enable_cpe_for_dv2_reg_vit(model: DinoWrapper,
119
+ max_img_size: Union[int, Tuple[int, int]] = 1024,
120
+ num_cls_tokens: int = 1,
121
+ pos_dropout: float = 0.1,
122
+ register_multiple: int = Optional[None],
123
+ num_registers: int = Optional[None],
124
+ ):
125
+ patch_size = model.patch_size
126
+ embed_dim = model.embed_dim
127
+ input_dims = model.inner.patch_embed.patches_resolution
128
+ normalize_patches = not isinstance(model.inner.patch_embed.norm, nn.Identity)
129
+ cls_token = True
130
+
131
+ max_img_size = int(round(max_img_size / patch_size) * patch_size)
132
+
133
+ patch_generator = ViTPatchGenerator(
134
+ patch_size=patch_size,
135
+ embed_dim=embed_dim,
136
+ input_dims=input_dims,
137
+ normalize_patches=normalize_patches,
138
+ cls_token=cls_token,
139
+ max_input_dims=max_img_size,
140
+ pos_dropout=pos_dropout,
141
+ num_cls_tokens=num_cls_tokens,
142
+ register_multiple=register_multiple,
143
+ num_registers=num_registers,
144
+ patch_bias=True,
145
+ )
146
+
147
+ inner = model.inner
148
+ inner.patch_generator = patch_generator
149
+ inner.patch_embed = None
150
+ inner.cls_token = None
151
+ inner.pos_embed = None
152
+ inner.register_tokens = None
153
+ inner.patch_size = patch_size
154
+
155
+ model.forward_features = MethodType(_forward_cpe_dinov2, model)
156
+ model.forward_intermediates = MethodType(_forward_intermediates_cpe_dinov2, model)
157
+
158
+
159
+ def enable_cpe(model: nn.Module,
160
+ *args,
161
+ **kwargs,
162
+ ):
163
+ if isinstance(model, VisionTransformer):
164
+ _enable_cpe_for_timm_vit(model, *args, **kwargs)
165
+ elif isinstance(model, DinoWrapper):
166
+ _enable_cpe_for_dv2_reg_vit(model, *args, **kwargs)
167
+ elif isinstance(model, HybridModel):
168
+ _enable_cpe_for_timm_vit(model.vit, *args, **kwargs)
169
+ else:
170
+ raise ValueError(f'CPE not supported for this model type: {type(model)}')
enable_spectral_reparam.py CHANGED
@@ -1,7 +1,15 @@
 
 
 
 
 
 
 
 
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
@@ -25,7 +33,7 @@ class _SNReweight(_SpectralNorm):
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
 
@@ -45,14 +53,16 @@ class _SNReweight(_SpectralNorm):
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)
@@ -90,21 +100,20 @@ class _SNReweight(_SpectralNorm):
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)
@@ -114,59 +123,100 @@ class _AttnSNReweight(nn.Module):
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
 
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 logging import getLogger
10
  import math
11
  import os
12
+ from typing import Dict, List, Optional, Union, Tuple
13
  from types import MethodType
14
 
15
  import torch
 
33
 
34
  if init_norm_to_current:
35
  # This will set the numerator to match the denominator, which should preserve the original values
36
+ init_scale = self._get_sigma(weight, n_power_iterations=20).item()
37
  else:
38
  init_scale = 1.0
39
 
 
53
  self.scale = nn.Parameter(torch.tensor([[init_value]], dtype=torch.float32, device=weight.device))
54
 
55
  # Re-implementing this because we need to make division by sigma safe
56
+ def _get_sigma(self, weight: torch.Tensor, n_power_iterations: int = None) -> torch.Tensor:
57
+ if not n_power_iterations:
58
+ n_power_iterations = self.n_power_iterations
59
  if weight.ndim == 1:
60
  # Faster and more exact path, no need to approximate anything
61
  sigma = weight.norm()
62
  else:
63
  weight_mat = self._reshape_weight_to_matrix(weight)
64
  if self.training:
65
+ self._power_method(weight_mat, n_power_iterations)
66
  # See above on why we need to clone
67
  u = self._u.clone(memory_format=torch.contiguous_format)
68
  v = self._v.clone(memory_format=torch.contiguous_format)
 
100
  return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
101
 
102
 
103
+ class _ChunkedSNReweight(nn.Module):
104
+ def __init__(self, weight: torch.Tensor, num_chunks: int, *args, init_norm_to_current: bool = False, **kwargs):
105
  super().__init__()
106
 
107
+ self.num_chunks = num_chunks
108
+ parts = weight.split(weight.shape[0] // num_chunks, dim=0)
 
109
 
110
  self.parts = nn.ModuleList([
111
+ _SNReweight(p, *args, init_norm_to_current=init_norm_to_current, **kwargs)
112
+ for p in parts
113
  ])
114
 
115
  def forward(self, weight: torch.Tensor, *args, **kwargs):
116
+ parts = weight.split(weight.shape[0] // self.num_chunks, dim=0)
117
 
118
  parts = [
119
  fn(p)
 
123
  return torch.cat(parts, dim=0)
124
 
125
 
126
+ class _AttnSNReweight(_ChunkedSNReweight):
127
+ def __init__(self, weight: torch.Tensor, *args, init_norm_to_current: bool = False, renorm_values: bool = False, **kwargs):
128
+ super().__init__(weight, 3, *args, init_norm_to_current=init_norm_to_current, **kwargs)
129
+
130
+ if not renorm_values:
131
+ self.parts[2] = nn.Identity()
132
+
133
+
134
+ def enable_spectral_reparam(model: Union[nn.Module, List[nn.Module]],
135
  n_power_iterations: int = 1,
136
  eps: float = 1e-6,
137
  init_norm_to_current: bool = False,
138
  renorm_values: bool = True,
139
+ renorm_mlp: bool = True,
140
+ state_dict_guidance: Optional[Dict[str, torch.Tensor]] = None):
141
+ if isinstance(model, (list, tuple)):
142
+ for i, sub in enumerate(model):
143
+ sub_sd = state_dict_guidance[i] if isinstance(state_dict_guidance, (list, tuple)) else state_dict_guidance
144
+ enable_spectral_reparam(sub, n_power_iterations=n_power_iterations, eps=eps,
145
+ init_norm_to_current=init_norm_to_current, renorm_values=renorm_values,
146
+ renorm_mlp=renorm_mlp, state_dict_guidance=sub_sd)
147
+ return
148
+
149
+ print('Enabling spectral reparametrization')
150
+ args = dict(n_power_iterations=n_power_iterations, dim=0, eps=eps, init_norm_to_current=init_norm_to_current)
151
+ visited_prefixes = set()
152
+
153
+ def is_guidance_parametrized(name: str):
154
+ if state_dict_guidance is None:
155
+ return True
156
+
157
+ p_name = f'{name}.parametrizations'
158
+ is_prm = any(k for k in state_dict_guidance if k.startswith(p_name) and k.endswith('_sn_version'))
159
+ return is_prm
160
+
161
+ def parametrize_linear(linear: nn.Linear):
162
+ parametrize.register_parametrization(
163
+ linear,
164
+ 'weight',
165
+ _SNReweight(linear.weight, **args)
166
+ )
167
 
168
+ for name, mod in model.named_modules():
169
+ pref = '.'.join(name.split('.')[:-1])
170
+ if pref in visited_prefixes:
171
+ continue
172
+
173
+ if isinstance(mod, Attention) or name.endswith('.attn'):
174
+ if is_guidance_parametrized(f'{name}.qkv'):
175
+ parametrize.register_parametrization(
176
+ mod.qkv,
177
+ 'weight',
178
+ _AttnSNReweight(mod.qkv.weight, renorm_values=renorm_values, **args),
179
+ )
180
+ if hasattr(mod, 'proj') and is_guidance_parametrized(f'{name}.proj'):
181
+ parametrize_linear(mod.proj)
182
+ visited_prefixes.add(name)
183
+ elif name.endswith('mlp') and renorm_mlp and hasattr(mod, 'w12'):
184
+ if is_guidance_parametrized(f'{name}.w12'):
185
+ parametrize.register_parametrization(
186
+ mod.w12,
187
+ 'weight',
188
+ _ChunkedSNReweight(mod.w12.weight, num_chunks=2, **args),
189
+ )
190
+ if is_guidance_parametrized(f'{name}.w3'):
191
+ parametrize_linear(mod.w3)
192
+ visited_prefixes.add(name)
193
+ elif isinstance(mod, nn.Linear) and 'patch_generator' not in name and is_guidance_parametrized(name):
194
+ parametrize_linear(mod)
195
+
196
+
197
+ def configure_spectral_reparam_from_args(model: nn.Module, args, state_dict_guidance: Optional[Dict[str, torch.Tensor]] = None):
198
  spectral_reparam = getattr(args, 'spectral_reparam', False)
199
  if isinstance(spectral_reparam, bool) and spectral_reparam:
200
+ enable_spectral_reparam(model, init_norm_to_current=True, state_dict_guidance=state_dict_guidance)
201
  elif isinstance(spectral_reparam, dict):
202
  enable_spectral_reparam(
203
  model,
204
  n_power_iterations=spectral_reparam.get('n_power_iterations', 1),
205
  eps=spectral_reparam.get('eps', 1e-12),
206
+ init_norm_to_current=True,
207
+ state_dict_guidance=state_dict_guidance,
208
  )
209
 
210
 
211
  def disable_spectral_reparam(model: nn.Module):
212
+ print('Disabling spectral reparametrization')
213
+ for name, mod in model.named_modules():
214
+ if parametrize.is_parametrized(mod):
215
+ parametrize.remove_parametrizations(mod, 'weight')
 
 
 
216
  pass
217
 
218
 
219
+
220
  if __name__ == '__main__':
221
  import argparse
222
  from . import radio_model as create_model
extra_models.py ADDED
@@ -0,0 +1,206 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from distutils.version import LooseVersion
2
+ from types import MethodType
3
+ from typing import List, Optional, Tuple, Union
4
+ import warnings
5
+
6
+ import torch
7
+ from torch import nn
8
+ import torch.nn.functional as F
9
+
10
+ from timm.models.registry import register_model
11
+ from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
12
+
13
+ from .forward_intermediates import forward_intermediates
14
+ from .input_conditioner import InputConditioner
15
+
16
+ _has_torch_sdpa = hasattr(F, 'scaled_dot_product_attention')
17
+
18
+
19
+ class PaliGemmaWrapper(nn.Module):
20
+ def __init__(self, vis_model: nn.Module, embed_dim: int):
21
+ super().__init__()
22
+
23
+ self.vis_model = vis_model
24
+ self.embed_dim = embed_dim
25
+
26
+ @property
27
+ def patch_size(self):
28
+ return self.vis_model.embeddings.patch_size
29
+
30
+ @property
31
+ def blocks(self):
32
+ return self.vis_model.encoder.layers
33
+
34
+ @property
35
+ def embed_dim(self):
36
+ return self.vis_model.embeddings.embed_dim
37
+
38
+ def forward(self, x: torch.Tensor):
39
+ outputs = self.vis_model(
40
+ x,
41
+ return_dict=False,
42
+ interpolate_pos_encoding=True,
43
+ )
44
+
45
+ features = outputs[0].to(torch.float32)
46
+
47
+ summary = features.mean(dim=1)
48
+
49
+ return summary, features
50
+
51
+ def forward_features(self, x: torch.Tensor):
52
+ return self(x)
53
+
54
+
55
+ def _get_paligemma_model(repo: str, embed_dim: int = None, dtype: torch.dtype = torch.bfloat16):
56
+ from transformers import PaliGemmaForConditionalGeneration, __version__ as tx_version
57
+
58
+ if LooseVersion(tx_version) > LooseVersion('4.44.2'):
59
+ warnings.warn(f'Your transformers version "{tx_version}" is higher than 4.44.2, and for whatever reason, PaliGemma might be broken.')
60
+
61
+ extra_args = dict()
62
+
63
+ if dtype is not None:
64
+ extra_args['torch_dtype'] = dtype
65
+ rev = str(dtype).split('.')[-1]
66
+ extra_args['revision'] = rev
67
+
68
+ model = PaliGemmaForConditionalGeneration.from_pretrained(repo, **extra_args)
69
+
70
+ vis_model = model.vision_tower.vision_model
71
+
72
+ vis_model = PaliGemmaWrapper(vis_model, embed_dim)
73
+
74
+ return vis_model
75
+
76
+ @register_model
77
+ def paligemma_896_student(**kwargs):
78
+ model = _get_paligemma_model('google/paligemma-3b-pt-896', embed_dim=1152, dtype=None)
79
+
80
+ return model
81
+
82
+
83
+ def dv2_sdpa(self, x: torch.Tensor) -> torch.Tensor:
84
+ B, N, C = x.shape
85
+ qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
86
+
87
+ q, k, v = qkv[0], qkv[1], qkv[2]
88
+ x = F.scaled_dot_product_attention(
89
+ q, k, v,
90
+ is_causal=False,
91
+ dropout_p=self.attn_drop.p if self.training else 0.,
92
+ scale=self.scale,
93
+ )
94
+ x = x.transpose(1, 2).reshape(B, N, C)
95
+ x = self.proj(x)
96
+ x = self.proj_drop(x)
97
+ return x
98
+
99
+ def _load_dino_v2(dino_v2_model, cache_dir: Optional[str] = None, pretrained=True, **kwargs):
100
+ if cache_dir:
101
+ torch.hub.set_dir(cache_dir)
102
+ model: nn.Module = torch.hub.load(
103
+ 'facebookresearch/dinov2',
104
+ dino_v2_model,
105
+ pretrained=pretrained,
106
+ # **kwargs,
107
+ )
108
+
109
+ if _has_torch_sdpa:
110
+ for n, m in model.named_modules():
111
+ if n.endswith('.attn'):
112
+ m.forward = MethodType(dv2_sdpa, m)
113
+
114
+ return model
115
+
116
+ class DinoWrapper(nn.Module):
117
+ def __init__(self, dino_model: nn.Module):
118
+ super().__init__()
119
+
120
+ self.inner = dino_model
121
+ dino_model.blocks = nn.Sequential(*dino_model.blocks)
122
+
123
+ @property
124
+ def embed_dim(self):
125
+ return self.inner.embed_dim
126
+
127
+ @property
128
+ def patch_size(self):
129
+ return self.inner.patch_size
130
+
131
+ @property
132
+ def num_cls_tokens(self):
133
+ return getattr(self.inner, 'num_tokens', 1)
134
+
135
+ @property
136
+ def num_registers(self):
137
+ return getattr(self.inner, 'num_register_tokens', 0)
138
+
139
+ @property
140
+ def num_summary_tokens(self):
141
+ return self.num_cls_tokens + self.num_registers
142
+
143
+ @property
144
+ def blocks(self):
145
+ return self.inner.blocks
146
+
147
+ def forward(self, *args, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]:
148
+ parts = self.inner.forward_features(*args, **kwargs)
149
+
150
+ cls_token = parts['x_norm_clstoken']
151
+ features = parts['x_norm_patchtokens']
152
+
153
+ return cls_token, features
154
+
155
+ def forward_features(self, x: torch.Tensor):
156
+ x = self.inner.prepare_tokens_with_masks(x)
157
+ x = self.inner.blocks(x)
158
+ x_norm = self.inner.norm(x)
159
+
160
+ return x_norm[:, 0], x_norm[:, self.num_summary_tokens:]
161
+
162
+ def patchify(self, x: torch.Tensor) -> torch.Tensor:
163
+ return self.inner.prepare_tokens_with_masks(x)
164
+
165
+ def forward_intermediates(self,
166
+ x: torch.Tensor,
167
+ norm: bool = False,
168
+ **kwargs,
169
+ ) -> Union[List[torch.Tensor], Tuple[torch.Tensor, List[torch.Tensor]]]:
170
+ return forward_intermediates(
171
+ self,
172
+ patch_extractor=self.inner.prepare_tokens_with_masks,
173
+ num_summary_tokens=self.num_summary_tokens,
174
+ num_cls_tokens=self.num_cls_tokens,
175
+ norm=self.inner.norm if norm else lambda y: y,
176
+ x=x,
177
+ **kwargs,
178
+ )
179
+
180
+
181
+ def _dino_student(arch: str, **kwargs):
182
+ from . import dinov2_arch
183
+
184
+ factory = getattr(dinov2_arch, arch)
185
+ model = factory()
186
+
187
+ model = DinoWrapper(model)
188
+
189
+ conditioner = InputConditioner(
190
+ input_scale=1.0,
191
+ norm_mean=IMAGENET_DEFAULT_MEAN,
192
+ norm_std=IMAGENET_DEFAULT_STD,
193
+ )
194
+
195
+ model.input_conditioner = conditioner
196
+
197
+ return model
198
+
199
+
200
+ @register_model
201
+ def dino_v2_l_student(**kwargs):
202
+ return _dino_student('dinov2_vitl14_reg', **kwargs)
203
+
204
+ @register_model
205
+ def dino_v2_g_student(**kwargs):
206
+ return _dino_student('dinov2_vitg14_reg', **kwargs)
extra_timm_models.py CHANGED
@@ -6,10 +6,24 @@
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
@@ -40,6 +54,34 @@ def vit_base_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:
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).
@@ -47,7 +89,7 @@ def vit_huge_patch16_224(pretrained=False, **kwargs) -> VisionTransformer:
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
@@ -64,3 +106,101 @@ def vit_huge_patch16_224_mlpnorm(pretrained=False, **kwargs) -> VisionTransforme
64
  m.norm = nn.LayerNorm(m.fc1.out_features)
65
 
66
  return model
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ import warnings
11
+
12
+ import torch
13
  from torch import nn
14
+ from torch.nn import functional as F
15
 
16
  from timm.models import register_model
17
+ from timm.models.vision_transformer import (
18
+ VisionTransformer,
19
+ _create_vision_transformer as _timm_create_vision_transformer,
20
+ Mlp,
21
+ Block,
22
+ LayerScale as TIMMLayerScale,
23
+ )
24
+
25
+ # Import these to also register them
26
+ from . import dinov2_arch
27
 
28
 
29
  @register_model
 
54
  return model
55
 
56
 
57
+ @register_model
58
+ def vit_base_patch16_v2_224(pretrained=False, **kwargs) -> VisionTransformer:
59
+ """ ViT-Base (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).
60
+ ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.
61
+ """
62
+ model_args = dict(
63
+ patch_size=16, embed_dim=768, depth=12, num_heads=12, init_values=1e-5,
64
+ reg_tokens=4, no_embed_class=True, img_size=518 * 16 // 14
65
+ )
66
+ model = _create_vision_transformer(
67
+ 'vit_base_patch14_reg4_dinov2', pretrained=pretrained, **dict(model_args, **kwargs))
68
+ return model
69
+
70
+
71
+ @register_model
72
+ def vit_large_patch16_v2_224(pretrained: bool = False, **kwargs) -> VisionTransformer:
73
+ """ ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).
74
+ ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.
75
+ """
76
+ name = 'vit_large_patch14_reg4_dinov2'
77
+ model_args = dict(
78
+ patch_size=16, embed_dim=1024, depth=24, num_heads=16, init_values=1e-5,
79
+ reg_tokens=4, no_embed_class=True, img_size=518 * 16 // 14
80
+ )
81
+ model = _create_vision_transformer(name, pretrained=pretrained, **dict(model_args, **kwargs))
82
+
83
+ return model
84
+
85
  @register_model
86
  def vit_huge_patch16_224(pretrained=False, **kwargs) -> VisionTransformer:
87
  """ ViT-Huge model (ViT-H/16) from original paper (https://arxiv.org/abs/2010.11929).
 
89
  model_args = dict(patch_size=16, embed_dim=1280, depth=32, num_heads=16)
90
  if pretrained:
91
  # There is no pretrained version of ViT-H/16, but we can adapt a ViT-H/14 for this purpose
92
+ model = _create_vision_transformer('vit_huge_patch14_224', pretrained=True, **dict(model_args, **kwargs))
93
  else:
94
  model = _create_vision_transformer('vit_huge_patch16_224', pretrained=False, **dict(model_args, **kwargs))
95
  return model
 
106
  m.norm = nn.LayerNorm(m.fc1.out_features)
107
 
108
  return model
109
+
110
+
111
+ @register_model
112
+ def vit_giant_patch16_224(pretrained=False, scaled_ln: bool = False, **kwargs) -> VisionTransformer:
113
+ """ ViT-giant model (ViT-g/16) from original paper (https://arxiv.org/abs/2010.11929).
114
+ """
115
+ model_args = dict(patch_size=16, embed_dim=1536, depth=40, num_heads=24)
116
+ model = _create_vision_transformer('vit_giant_patch16_224', pretrained=False, **dict(model_args, **kwargs))
117
+ if scaled_ln:
118
+ _apply_scaled_ln(model)
119
+ return model
120
+
121
+
122
+ @register_model
123
+ def vit_bigG_patch14_224(pretrained=False, **kwargs) -> VisionTransformer:
124
+ model_args = dict(patch_size=14, embed_dim=1664, depth=48, num_heads=16, init_values=1e-6)
125
+ model = _create_vision_transformer('vit_bigG_patch14', pretrained=False, **dict(model_args, **kwargs))
126
+ return model
127
+
128
+
129
+ def _create_vision_transformer(*args, **kwargs):
130
+ model = _timm_create_vision_transformer(*args, **kwargs)
131
+ _patch_layer_scale(model)
132
+ return model
133
+
134
+
135
+ def _patch_layer_scale(model: VisionTransformer):
136
+ def replace_ls(old_ls: TIMMLayerScale):
137
+ new_ls = dinov2_arch.LayerScale(old_ls.gamma.shape[0], inplace=old_ls.inplace)
138
+ new_ls.load_state_dict(old_ls.state_dict())
139
+ return new_ls
140
+
141
+ # Monkey patch: Replace TIMM's LayerScale with our modified DINOv2 one, that uses a param name
142
+ # other than gamma, so that HFHub doesn't mess with it!
143
+ for mod in model.modules():
144
+ if isinstance(mod, Block):
145
+ if isinstance(mod.ls1, TIMMLayerScale):
146
+ mod.ls1 = replace_ls(mod.ls1)
147
+ if isinstance(mod.ls2, TIMMLayerScale):
148
+ mod.ls2 = replace_ls(mod.ls2)
149
+ pass
150
+
151
+
152
+ class ScaledLayerNorm(nn.LayerNorm):
153
+ '''
154
+ https://arxiv.org/pdf/2502.05795v1
155
+ '''
156
+ def __init__(self, ln_base: nn.LayerNorm, depth: int = 0):
157
+ super().__init__(ln_base.normalized_shape, eps=ln_base.eps, elementwise_affine=ln_base.elementwise_affine)
158
+ self.load_state_dict(ln_base.state_dict())
159
+ self.register_buffer('ln_scale', torch.tensor(1.0 / math.sqrt(depth)), persistent=False)
160
+
161
+ def forward(self, x):
162
+ y = super().forward(x)
163
+ y = y * self.ln_scale
164
+ return y
165
+
166
+
167
+ class DyT(nn.Module):
168
+ def __init__(self, C: int, init_alpha: float):
169
+ super().__init__()
170
+ self.alpha = nn.Parameter(torch.full((1,), init_alpha))
171
+ self.gamma = nn.Parameter(torch.ones(C))
172
+ self.beta = nn.Parameter(torch.zeros(C))
173
+
174
+ def forward(self, x: torch.Tensor):
175
+ x = F.tanh(self.alpha * x)
176
+ return self.gamma * x + self.beta
177
+
178
+ @register_model
179
+ def vit_large_dyt_patch16_224(pretrained: bool = False, **kwargs) -> VisionTransformer:
180
+ """ ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).
181
+ ImageNet-1k weights fine-tuned from in21k @ 224x224, source https://github.com/google-research/vision_transformer.
182
+ """
183
+ model_args = dict(patch_size=16, embed_dim=1024, depth=24, num_heads=16)
184
+ model = _create_vision_transformer('vit_large_dyt_patch16_224', pretrained=pretrained, **dict(model_args, **kwargs))
185
+
186
+ def _replace_ln_with_dyt(ln: nn.LayerNorm, depth: int):
187
+ return DyT(ln.normalized_shape[0], init_alpha=0.9)
188
+ _replace_ln(model, _replace_ln_with_dyt)
189
+
190
+ return model
191
+
192
+
193
+ def _apply_scaled_ln(model: VisionTransformer):
194
+ warnings.warn('Post-LayerNorm scaling activated!')
195
+
196
+ _replace_ln(model, lambda ln, depth: ScaledLayerNorm(ln, depth=depth))
197
+
198
+ def _replace_ln(model: VisionTransformer, fn):
199
+ def _inner_replace_ln(block: Block, depth: int, key: str):
200
+ prev = getattr(block, key)
201
+ if isinstance(prev, nn.LayerNorm):
202
+ setattr(block, key, fn(prev, depth=depth))
203
+
204
+ for i, block in enumerate(model.blocks):
205
+ _inner_replace_ln(block, i + 1, 'norm1')
206
+ _inner_replace_ln(block, i + 1, 'norm2')
feature_normalizer.py ADDED
@@ -0,0 +1,111 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 collections import namedtuple
9
+ from typing import NamedTuple, Optional, Tuple
10
+ import torch
11
+ from torch import nn
12
+
13
+
14
+ def _run_kernel(x: torch.Tensor, mean: torch.Tensor, tx: torch.Tensor):
15
+ if x.ndim <= 3:
16
+ x = x - mean
17
+ x = x @ tx.T
18
+ elif x.ndim == 4:
19
+ x = x - mean.reshape(1, -1, 1, 1)
20
+ kernel = tx.reshape(*tx.shape, 1, 1)
21
+ x = torch.nn.functional.conv2d(x, weight=kernel, bias=None, stride=1, padding=0)
22
+ else:
23
+ raise ValueError(f'Unsupported input dimension: {x.ndim}, shape: {x.shape}')
24
+ return x
25
+
26
+
27
+ class FeatureNormalizer(nn.Module):
28
+ def __init__(self, embed_dim: int, dtype: torch.dtype = torch.float32):
29
+ super().__init__()
30
+
31
+ self.register_buffer('mean', torch.zeros(embed_dim, dtype=dtype))
32
+ self.register_buffer('tx', torch.eye(embed_dim, dtype=dtype))
33
+
34
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
35
+ x = _run_kernel(x, self.mean, self.tx)
36
+ return x
37
+
38
+
39
+ class InterFeatState(NamedTuple):
40
+ y: torch.Tensor
41
+ alpha: torch.Tensor
42
+
43
+
44
+ class IntermediateFeatureNormalizerBase(nn.Module):
45
+ def forward(self, x: torch.Tensor, index: int, rot_index: int = None, skip: Optional[int] = None) -> InterFeatState:
46
+ raise NotImplementedError()
47
+
48
+
49
+ class IntermediateFeatureNormalizer(IntermediateFeatureNormalizerBase):
50
+ def __init__(self, num_intermediates: int, embed_dim: int, rot_per_layer: bool = False, dtype: torch.dtype = torch.float32):
51
+ super().__init__()
52
+ self.register_buffer('alphas', torch.ones(num_intermediates, dtype=dtype))
53
+
54
+ rot = torch.eye(embed_dim, dtype=dtype)
55
+ if rot_per_layer:
56
+ rot = rot.unsqueeze(0).repeat(num_intermediates, 1, 1)
57
+
58
+ self.register_buffer('rotation', rot.contiguous())
59
+ self.register_buffer('means', torch.zeros(num_intermediates, embed_dim, dtype=dtype))
60
+
61
+ def forward(self, x: torch.Tensor, index: int, rot_index: int = None, skip: Optional[int] = None) -> InterFeatState:
62
+ if rot_index is None:
63
+ rot_index = index
64
+
65
+ if skip:
66
+ assert x.ndim == 3, f'Cannot use the `skip` parameter when the `x` tensor isn\'t 3-dimensional.'
67
+ prefix, x = x[:, :skip], x[:, skip:]
68
+
69
+ rotation = self._get_rotation(rot_index)
70
+ y = _run_kernel(x, self.means[index], rotation)
71
+
72
+ alpha = self.alphas[index]
73
+ if skip:
74
+ alpha = torch.cat([
75
+ torch.ones(skip, dtype=alpha.dtype, device=alpha.device),
76
+ alpha[None].expand(y.shape[1]),
77
+ ]).reshape(1, -1, 1)
78
+ y = torch.cat([prefix, y], dim=1)
79
+ else:
80
+ if x.ndim == 3:
81
+ alpha = alpha.reshape(1, 1, 1).expand(1, y.shape[1], 1)
82
+ elif x.ndim == 4:
83
+ alpha = alpha.reshape(1, 1, 1, 1).expand(1, 1, *y.shape[2:])
84
+ else:
85
+ raise ValueError(f'Unsupported input dimension: {x.ndim}')
86
+
87
+ return InterFeatState(y, alpha)
88
+
89
+ def _get_rotation(self, rot_index: int) -> torch.Tensor:
90
+ if self.rotation.ndim == 2:
91
+ return self.rotation
92
+ return self.rotation[rot_index]
93
+
94
+
95
+ class NullIntermediateFeatureNormalizer(IntermediateFeatureNormalizerBase):
96
+ instances = dict()
97
+
98
+ def __init__(self, dtype: torch.dtype, device: torch.device):
99
+ super().__init__()
100
+ self.register_buffer('alpha', torch.tensor(1, dtype=dtype, device=device))
101
+
102
+ @staticmethod
103
+ def get_instance(dtype: torch.dtype, device: torch.device):
104
+ instance = NullIntermediateFeatureNormalizer.instances.get((dtype, device), None)
105
+ if instance is None:
106
+ instance = NullIntermediateFeatureNormalizer(dtype, device)
107
+ NullIntermediateFeatureNormalizer.instances[(dtype, device)] = instance
108
+ return instance
109
+
110
+ def forward(self, x: torch.Tensor, index: int, rot_index: int = None, skip: Optional[int] = None) -> InterFeatState:
111
+ return InterFeatState(x, self.alpha)
forward_intermediates.py ADDED
@@ -0,0 +1,138 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 Callable, Dict, List, Optional, Set, Tuple, Union, Any, Iterable
10
+ from types import MethodType
11
+
12
+ import torch
13
+ from torch import nn
14
+
15
+ from .feature_normalizer import IntermediateFeatureNormalizerBase, NullIntermediateFeatureNormalizer
16
+
17
+
18
+ def _take_indices(
19
+ num_blocks: int,
20
+ n: Optional[Union[int, List[int], Tuple[int]]],
21
+ ) -> Tuple[Set[int], int]:
22
+ if isinstance(n, int):
23
+ assert n >= 0
24
+ take_indices = {x for x in range(num_blocks - n, num_blocks)}
25
+ else:
26
+ take_indices = {num_blocks + idx if idx < 0 else idx for idx in n}
27
+ return take_indices, max(take_indices)
28
+
29
+
30
+ def forward_intermediates(
31
+ model: nn.Module,
32
+ patch_extractor: Callable[[torch.Tensor], torch.Tensor],
33
+ norm: nn.Module,
34
+ num_summary_tokens: int,
35
+ num_cls_tokens: int,
36
+ x: torch.Tensor,
37
+ indices: Optional[Union[int, List[int], Tuple[int]]] = None,
38
+ return_prefix_tokens: bool = False,
39
+ stop_early: bool = False,
40
+ output_fmt: str = 'NCHW',
41
+ intermediates_only: bool = False,
42
+ aggregation: Optional[str] = "sparse",
43
+ inter_feature_normalizer: Optional[IntermediateFeatureNormalizerBase] = None,
44
+ norm_alpha_scheme = "post-alpha",
45
+ block_kwargs: Dict = None,
46
+ ) -> Union[List[torch.Tensor], Tuple[torch.Tensor, List[torch.Tensor]]]:
47
+ """ Forward features that returns intermediates.
48
+
49
+ The Dense layer aggregation method is inspired from the paper: "Dense Connector for MLLMs"
50
+ by Yao, Huanjin et al. (2024). arXiv preprint arXiv:2405.13800}
51
+
52
+ Args:
53
+ x: Input image tensor
54
+ indices: Take last n blocks if int, select matching indices if sequence
55
+ return_prefix_tokens: Return both prefix and spatial intermediate tokens
56
+ norm: Apply norm layer to all intermediates
57
+ stop_early: Stop iterating over blocks when last desired intermediate hit
58
+ output_fmt: Shape of intermediate feature outputs
59
+ intermediates_only: Only return intermediate features
60
+ aggregation: intermediate layer aggregation method (sparse or dense)
61
+ norm_alpha_scheme: apply alpha before ("pre-alpha") or after accumulation ("post-alpha")
62
+ Returns:
63
+ """
64
+ assert output_fmt in ('NCHW', 'NLC'), 'Output format must be one of NCHW or NLC.'
65
+ assert aggregation in ('sparse', 'dense'), 'Aggregation must be one of sparse or dense.'
66
+ reshape = output_fmt == 'NCHW'
67
+ intermediates = []
68
+
69
+ block_kwargs = block_kwargs or dict()
70
+
71
+ blocks = model.blocks
72
+
73
+ take_indices, max_index = _take_indices(len(blocks), indices)
74
+ take_indices = sorted(take_indices)
75
+ # forward pass
76
+ B, _, height, width = x.shape
77
+
78
+ x = patch_extractor(x)
79
+
80
+ if stop_early:
81
+ blocks = blocks[:max_index + 1]
82
+
83
+ if inter_feature_normalizer is None or norm_alpha_scheme == 'none':
84
+ inter_feature_normalizer = NullIntermediateFeatureNormalizer.get_instance(x.dtype, x.device)
85
+
86
+ assert norm_alpha_scheme in ('none', 'pre-alpha', 'post-alpha'), f'Unsupported alpha scheme: {norm_alpha_scheme}'
87
+ post_alpha_scheme = norm_alpha_scheme == 'post-alpha'
88
+
89
+ accumulator = 0
90
+ alpha_sum = 0
91
+ num_accumulated = 0
92
+
93
+ take_off = 0
94
+
95
+ for i, blk in enumerate(blocks):
96
+ x = blk(x, **block_kwargs)
97
+ if aggregation == "dense":
98
+ # Arbitrarily use the rotation matrix from the final layer in the dense group
99
+ y, alpha = inter_feature_normalizer(x, i, rot_index=take_indices[take_off], skip=num_summary_tokens)
100
+ if post_alpha_scheme:
101
+ accumulator = accumulator + y
102
+ alpha_sum = alpha_sum + alpha
103
+ else:
104
+ accumulator = accumulator + (alpha * y)
105
+ alpha_sum += 1
106
+ num_accumulated += 1
107
+ if i == take_indices[take_off]:
108
+ if aggregation == "dense":
109
+ alpha = alpha_sum / num_accumulated
110
+ x_ = alpha * accumulator / num_accumulated
111
+ num_accumulated = 0
112
+ accumulator = 0
113
+ alpha_sum = 0
114
+ else:
115
+ y, alpha = inter_feature_normalizer(x, i, skip=num_summary_tokens)
116
+ x_ = alpha * y
117
+ # normalize intermediates with final norm layer if enabled
118
+ intermediates.append(norm(x_))
119
+ take_off = min(take_off + 1, len(take_indices) - 1)
120
+
121
+ # process intermediates
122
+
123
+ # split prefix (e.g. class, distill) and spatial feature tokens
124
+ prefix_tokens = [y[:, :num_cls_tokens] for y in intermediates]
125
+ intermediates = [y[:, num_summary_tokens:] for y in intermediates]
126
+
127
+ if reshape:
128
+ # reshape to BCHW output format
129
+ H = height // model.patch_size
130
+ W = width // model.patch_size
131
+ intermediates = [y.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() for y in intermediates]
132
+ if not torch.jit.is_scripting() and return_prefix_tokens:
133
+ # return_prefix not support in torchscript due to poor type handling
134
+ intermediates = list(zip(prefix_tokens, intermediates))
135
+ if intermediates_only:
136
+ return intermediates
137
+ x = norm(x)
138
+ return x, intermediates
hf_model.py CHANGED
@@ -12,7 +12,7 @@
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
@@ -25,12 +25,15 @@ from .common import RESOURCE_MAP, DEFAULT_VERSION
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
@@ -40,6 +43,7 @@ from .vitdet import apply_vitdet_arch, VitDetArgs
40
 
41
  # Register extra models
42
  from .extra_timm_models import *
 
43
 
44
 
45
  class RADIOConfig(PretrainedConfig):
@@ -53,7 +57,10 @@ class RADIOConfig(PretrainedConfig):
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
@@ -71,10 +78,14 @@ class RADIOConfig(PretrainedConfig):
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
 
@@ -106,13 +117,28 @@ class RADIOModel(PreTrainedModel):
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,
@@ -123,6 +149,8 @@ class RADIOModel(PreTrainedModel):
123
  window_size=config.vitdet_window_size,
124
  preferred_resolution=config.preferred_resolution,
125
  adaptors=adaptors,
 
 
126
  )
127
 
128
  @property
 
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, Dict, Optional, List, Union
16
 
17
  from timm.models import VisionTransformer
18
  import torch
 
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_config
29
  from .adaptor_registry import adaptor_registry
30
  from .cls_token import ClsToken
31
+ from .dinov2_arch import dinov2_vitg14_reg
32
  from .enable_cpe_support import enable_cpe
33
  from .enable_spectral_reparam import configure_spectral_reparam_from_args
34
  from .eradio_model import eradio
35
+ from .feature_normalizer import FeatureNormalizer, IntermediateFeatureNormalizer
36
+ from .forward_intermediates import forward_intermediates
37
  from .radio_model import create_model_from_args
38
  from .radio_model import RADIOModel as RADIOModelBase, Resolution
39
  from .input_conditioner import get_default_conditioner, InputConditioner
 
43
 
44
  # Register extra models
45
  from .extra_timm_models import *
46
+ from .extra_models import *
47
 
48
 
49
  class RADIOConfig(PretrainedConfig):
 
57
  max_resolution: Optional[int] = None,
58
  preferred_resolution: Optional[Resolution] = None,
59
  adaptor_names: Union[str, List[str]] = None,
60
+ adaptor_configs: Dict[str, Dict[str, int]] = None,
61
  vitdet_window_size: Optional[int] = None,
62
+ feature_normalizer_config: Optional[dict] = None,
63
+ inter_feature_normalizer_config: Optional[dict] = None,
64
  **kwargs,
65
  ):
66
  self.args = args
 
78
  preferred_resolution or resource.preferred_resolution
79
  )
80
  self.adaptor_names = adaptor_names
81
+ self.adaptor_configs = adaptor_configs
82
  self.vitdet_window_size = vitdet_window_size
83
+ self.feature_normalizer_config = feature_normalizer_config
84
+ self.inter_feature_normalizer_config = inter_feature_normalizer_config
85
  super().__init__(**kwargs)
86
 
87
 
88
+
89
  class RADIOModel(PreTrainedModel):
90
  """Pretrained Hugging Face model for RADIO.
91
 
 
117
  dtype=torch.int64,
118
  )
119
 
120
+ adaptor_configs = config.adaptor_configs
121
+ adaptor_names = config.adaptor_names or []
 
 
 
122
 
123
  adaptors = dict()
124
+ for adaptor_name in adaptor_names:
125
+ mlp_config = adaptor_configs[adaptor_name]
126
+ adaptor = GenericAdaptor(args, None, None, mlp_config)
127
+ adaptor.head_idx = mlp_config["head_idx"]
128
+ adaptors[adaptor_name] = adaptor
129
+
130
+ feature_normalizer = None
131
+ if config.feature_normalizer_config is not None:
132
+ # Actual normalization values will be restored when loading checkpoint weights.
133
+ feature_normalizer = FeatureNormalizer(config.feature_normalizer_config["embed_dim"])
134
+
135
+ inter_feature_normalizer = None
136
+ if config.inter_feature_normalizer_config is not None:
137
+ inter_feature_normalizer = IntermediateFeatureNormalizer(
138
+ config.inter_feature_normalizer_config["num_intermediates"],
139
+ config.inter_feature_normalizer_config["embed_dim"],
140
+ rot_per_layer=config.inter_feature_normalizer_config["rot_per_layer"],
141
+ dtype=dtype)
142
 
143
  self.radio_model = RADIOModelBase(
144
  model,
 
149
  window_size=config.vitdet_window_size,
150
  preferred_resolution=config.preferred_resolution,
151
  adaptors=adaptors,
152
+ feature_normalizer=feature_normalizer,
153
+ inter_feature_normalizer=inter_feature_normalizer,
154
  )
155
 
156
  @property
model.safetensors CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:608c1132b7a82bf4652827a209509864b5d7d925aa2d2d256e2e1a2968b16792
3
- size 392950088
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d58db8f4ba0d3f6ea2dc8bc9bc846e6f059c656160f23e576c0b15368b8c4770
3
+ size 395312688
radio_model.py CHANGED
@@ -5,7 +5,7 @@
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
@@ -14,11 +14,11 @@ 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):
@@ -37,6 +37,8 @@ class RADIOModel(nn.Module):
37
  summary_idxs: Optional[torch.Tensor] = None,
38
  window_size: int = None,
39
  adaptors: Dict[str, AdaptorBase] = None,
 
 
40
  ):
41
  super().__init__()
42
 
@@ -55,12 +57,32 @@ class RADIOModel(nn.Module):
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
 
@@ -68,6 +90,8 @@ class RADIOModel(nn.Module):
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
@@ -92,6 +116,17 @@ class RADIOModel(nn.Module):
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()
@@ -111,7 +146,14 @@ class RADIOModel(nn.Module):
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`. '
@@ -120,7 +162,10 @@ class RADIOModel(nn.Module):
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:
@@ -147,18 +192,40 @@ class RADIOModel(nn.Module):
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
 
@@ -174,6 +241,7 @@ class RADIOModel(nn.Module):
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:
@@ -182,14 +250,17 @@ class RADIOModel(nn.Module):
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,
@@ -198,12 +269,33 @@ class RADIOModel(nn.Module):
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:
@@ -238,20 +330,14 @@ def create_model_from_args(args) -> nn.Module:
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
 
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, Iterable, List, NamedTuple, Optional, Tuple, Union
9
 
10
  import torch
11
  from torch import nn
 
14
 
15
  from .enable_cpe_support import enable_cpe
16
  from .input_conditioner import InputConditioner
 
 
17
  from .adaptor_base import AdaptorBase, RadioOutput, AdaptorInput
18
  from . import eradio_model
19
  from .enable_spectral_reparam import configure_spectral_reparam_from_args
20
+ from .feature_normalizer import FeatureNormalizer, IntermediateFeatureNormalizer
21
+ from . import dual_hybrid_vit
22
 
23
 
24
  class Resolution(NamedTuple):
 
37
  summary_idxs: Optional[torch.Tensor] = None,
38
  window_size: int = None,
39
  adaptors: Dict[str, AdaptorBase] = None,
40
+ feature_normalizer: Optional[FeatureNormalizer] = None,
41
+ inter_feature_normalizer: Optional[IntermediateFeatureNormalizer] = None,
42
  ):
43
  super().__init__()
44
 
 
57
  adaptors = adaptors or dict()
58
  self.adaptors = nn.ModuleDict(adaptors)
59
 
60
+ if feature_normalizer is None:
61
+ feature_normalizer = nn.Identity()
62
+ self.feature_normalizer = feature_normalizer
63
+ self.inter_feature_normalizer = inter_feature_normalizer
64
+
65
  @property
66
  def num_summary_tokens(self) -> int:
67
+ if hasattr(self.model, 'num_summary_tokens'):
68
+ return self.model.num_summary_tokens
69
+
70
  patch_gen = getattr(self.model, "patch_generator", None)
71
  if patch_gen is not None:
72
  return patch_gen.num_skip
73
+ elif getattr(self.model, 'global_pool', None) == 'avg':
74
+ return 0
75
+ return 1
76
+
77
+ @property
78
+ def num_cls_tokens(self) -> int:
79
+ if hasattr(self.model, 'num_cls_tokens'):
80
+ return self.model.num_cls_tokens
81
+
82
+ patch_gen = getattr(self.model, 'patch_generator', None)
83
+ if patch_gen is not None:
84
+ return patch_gen.num_cls_tokens
85
+ elif getattr(self.model, 'global_pool', None) == 'avg':
86
  return 0
87
  return 1
88
 
 
90
  def patch_size(self) -> int:
91
  if self._patch_size is not None:
92
  return self._patch_size
93
+ if hasattr(self.model, "patch_size"):
94
+ return self.model.patch_size
95
  patch_gen = getattr(self.model, "patch_generator", None)
96
  if patch_gen is not None:
97
  return patch_gen.patch_size
 
116
  res *= self.window_size
117
  return res
118
 
119
+ @property
120
+ def blocks(self) -> Iterable[nn.Module]:
121
+ blocks = getattr(self.model, 'blocks', None)
122
+ if blocks is not None:
123
+ return blocks
124
+ return None
125
+
126
+ @property
127
+ def embed_dim(self) -> int:
128
+ return self.model.embed_dim
129
+
130
  def make_preprocessor_external(self) -> Callable[[torch.Tensor], torch.Tensor]:
131
  ret = self.input_conditioner
132
  self.input_conditioner = nn.Identity()
 
146
  if fn is not None:
147
  fn()
148
 
149
+ def forward(self, x: torch.Tensor, feature_fmt: str = 'NLC') -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
150
+ '''
151
+ Forward process for model.
152
+ Args:
153
+ x: Input tensor. Unless `make_preprocessor_external` has been called, then the dynamic range of `x` is expected to be `[0, 1]`,
154
+ otherwise `x` is expected to be mean centered with unit standard deviation.
155
+ feature_format: ['NLC', 'NCHW'] - The output format for the features.
156
+ '''
157
  res_step = self.min_resolution_step
158
  if res_step is not None and (x.shape[-2] % res_step != 0 or x.shape[-1] % res_step != 0):
159
  raise ValueError('The input resolution must be a multiple of `self.min_resolution_step`. '
 
162
 
163
  x = self.input_conditioner(x)
164
  y = self.model.forward_features(x)
165
+ ret = self._extract_final(x, y, feature_fmt=feature_fmt)
166
+ return ret
167
 
168
+ def _extract_final(self, x: torch.Tensor, y: torch.Tensor, feature_fmt: str = 'NLC'):
169
  if isinstance(self.model, VisionTransformer):
170
  patch_gen = getattr(self.model, "patch_generator", None)
171
  if patch_gen is not None:
 
192
  all_summary, all_feat = y
193
  bb_summary = all_summary
194
  else:
195
+ all_summary = y[:, :self.num_cls_tokens]
196
+ if self.summary_idxs is not None and all_summary.shape[1] > 1:
197
+ if all_summary.shape[1] == 1:
198
+ # Create dummy duplicates
199
+ all_summary = all_summary.expand(-1, 128, -1)
200
+ bb_summary = all_summary[:, self.summary_idxs]
201
+ else:
202
+ bb_summary = all_summary
203
+ all_feat = y[:, self.num_summary_tokens:]
204
+
205
+ all_feat = self.feature_normalizer(all_feat)
206
+
207
+ if feature_fmt == 'NCHW':
208
+ fmt_feat = (all_feat.reshape(all_feat.shape[0], x.shape[-2] // self.patch_size, x.shape[-1] // self.patch_size, all_feat.shape[2])
209
+ .permute(0, 3, 1, 2)
210
+ )
211
+ elif feature_fmt == 'NLC':
212
+ fmt_feat = all_feat
213
+ else:
214
+ raise ValueError(f'Unsupported feature_fmt: {feature_fmt}. Must be one of ["NLC", "NCHW"]')
215
+
216
+ ret = RadioOutput(bb_summary.flatten(1), fmt_feat)
217
 
 
 
218
  if self.adaptors:
219
  ret = dict(backbone=ret)
220
  for name, adaptor in self.adaptors.items():
221
  if all_summary.ndim == 3:
222
+ if all_summary.shape[1] == 1:
223
+ summary = all_summary[:, 0]
224
+ else:
225
+ summary = all_summary[:, adaptor.head_idx]
226
  else:
227
  summary = all_summary
228
+ ada_input = AdaptorInput(images=x, summary=summary.float(), features=all_feat, feature_fmt=feature_fmt, patch_size=self.patch_size)
229
  v = adaptor(ada_input).to(torch.float32)
230
  ret[name] = v
231
 
 
241
  output_fmt: str = 'NCHW',
242
  intermediates_only: bool = False,
243
  aggregation: Optional[str] = "sparse",
244
+ norm_alpha_scheme: Optional[str] = "post-alpha",
245
  ) -> List[RadioOutput]:
246
  """ Forward features that returns intermediates.
247
  Args:
 
250
  return_prefix_tokens: Return both prefix and spatial intermediate tokens
251
  norm: Apply norm layer to all intermediates
252
  stop_early: Stop iterating over blocks when last desired intermediate hit
253
+ output_fmt: Shape of intermediate feature outputs. Options: NCHW, NLC
254
  intermediates_only: Only return intermediate features
255
  aggregation: intermediate layer aggregation method (sparse or dense).
256
  Dense accumulation is done by averaging the features in each group.
257
+ norm_alpha_scheme: apply alpha before ("pre-alpha") or after accumulation ("post-alpha"), or don't normalize ("none")
258
+ Only affects dense aggregation
259
  Returns:
260
  List of RadioOutput objects.
261
  """
262
+ x = self.input_conditioner(x)
263
+ intermediates = self.model.forward_intermediates(
264
  x,
265
  indices=indices,
266
  return_prefix_tokens=return_prefix_tokens,
 
269
  output_fmt=output_fmt,
270
  intermediates_only=intermediates_only,
271
  aggregation=aggregation,
272
+ inter_feature_normalizer=self.inter_feature_normalizer,
273
+ norm_alpha_scheme=norm_alpha_scheme,
274
  )
275
+
276
+ if not intermediates_only:
277
+ final, intermediates = intermediates
278
+
279
+ def prepare_summary(summ: Optional[torch.Tensor]):
280
+ if summ is None:
281
+ return summ
282
+ if self.summary_idxs is not None and summ.shape[1] > 1:
283
+ summ = summ[:, self.summary_idxs]
284
+ return summ.flatten(1)
285
+
286
  if return_prefix_tokens:
287
+ radio_outputs = [
288
+ RadioOutput(prepare_summary(summary), features)
289
+ for summary, features in intermediates
290
+ ]
291
  else:
292
+ radio_outputs = intermediates
293
+
294
+ if intermediates_only:
295
+ return radio_outputs
296
+ else:
297
+ final = self._extract_final(x, final, feature_fmt=output_fmt)
298
+ return final, radio_outputs
299
 
300
 
301
  def create_model_from_args(args) -> nn.Module:
 
330
 
331
  model.head = nn.Identity()
332
 
 
 
 
 
333
  if args.cpe_max_size is not None:
334
  uq_teachers = set(t['name'] for t in args.teachers)
335
  enable_cpe(
336
  model,
337
  args.cpe_max_size,
338
  num_cls_tokens=len(uq_teachers) if args.cls_token_per_teacher else 1,
339
+ register_multiple=getattr(args, 'register_multiple', None),
340
+ num_registers=getattr(args, 'cpe_num_registers', None),
341
  )
342
 
 
 
 
343
  return model
vit_patch_generator.py CHANGED
@@ -36,7 +36,9 @@ class ViTPatchGenerator(nn.Module):
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__()
@@ -71,7 +73,7 @@ class ViTPatchGenerator(nn.Module):
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
@@ -82,6 +84,7 @@ class ViTPatchGenerator(nn.Module):
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()
@@ -103,6 +106,10 @@ class ViTPatchGenerator(nn.Module):
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
@@ -116,10 +123,6 @@ class ViTPatchGenerator(nn.Module):
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]))
@@ -274,26 +277,11 @@ class Im2Patches(nn.Module):
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)
 
36
  pos_dropout: float = 0.0,
37
  return_pos_enc: bool = False,
38
  num_cls_tokens: int = 1,
39
+ register_multiple: Optional[int] = None,
40
+ num_registers: Optional[int] = None,
41
+ patch_bias: bool = False,
42
  device=None, dtype=None,
43
  ):
44
  super().__init__()
 
73
  self.max_input_dims = max_input_dims
74
 
75
  self.im_to_patches = Im2Patches(patch_size)
76
+ self.embedder = ViTPatchLinear(patch_size, embed_dim, bias=patch_bias, **factory)
77
 
78
  if abs_pos:
79
  scale = embed_dim ** -0.5
 
84
  num_tokens=num_cls_tokens,
85
  enabled=cls_token,
86
  register_multiple=register_multiple,
87
+ num_registers=num_registers,
88
  )
89
 
90
  self.patch_normalizer = nn.LayerNorm(embed_dim) if normalize_patches else nn.Identity()
 
106
  def num_cls_tokens(self):
107
  return self.cls_token.num_tokens
108
 
109
+ @property
110
+ def num_cls_patches(self):
111
+ return self.cls_token.num_patches
112
+
113
  @property
114
  def num_registers(self):
115
  return self.cls_token.num_registers
 
123
  'pos_embed',
124
  ]
125
 
 
 
 
 
126
  def _load_embed(self, src_embed: torch.Tensor, targ_embed: nn.Parameter):
127
  if src_embed.shape != targ_embed.shape:
128
  src_size = int(math.sqrt(src_embed.shape[1]))
 
277
 
278
 
279
  class ViTPatchLinear(nn.Linear):
280
+ def __init__(self, patch_size: int, embed_dim: int, bias: bool = False, **factory):
281
  super().__init__(
282
  3 * (patch_size ** 2),
283
  embed_dim,
284
+ bias=bias,
285
  **factory
286
  )
287
  self.patch_size = patch_size
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
vitdet.py CHANGED
@@ -12,6 +12,8 @@ from torch import nn
12
  from timm.models import VisionTransformer
13
  from einops import rearrange
14
 
 
 
15
  DEFAULT_NUM_WINDOWED = 5
16
  DEFAULT_NUM_GLOBAL = 4
17
 
@@ -29,11 +31,16 @@ class VitDetArgs:
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
 
 
12
  from timm.models import VisionTransformer
13
  from einops import rearrange
14
 
15
+ from .extra_models import DinoWrapper
16
+
17
  DEFAULT_NUM_WINDOWED = 5
18
  DEFAULT_NUM_GLOBAL = 4
19
 
 
31
  self.num_global = num_global
32
 
33
 
34
+ def apply_vitdet_arch(model: Union[VisionTransformer, DinoWrapper], args: VitDetArgs):
35
  if isinstance(model, VisionTransformer):
36
  patch_embed = getattr(model, 'patch_generator', model.patch_embed)
37
 
38
  return ViTDetHook(patch_embed, model.blocks, args)
39
+ elif isinstance(model, DinoWrapper):
40
+ inner = model.inner
41
+
42
+ patch_embed = getattr(inner, 'patch_generator', inner.patch_embed)
43
+ return ViTDetHook(patch_embed, inner.blocks, args)
44
  else:
45
  print(f'Warning: Unable to apply VitDet aug!', file=sys.stderr)
46