monai-weekly 1.5.dev2446__py3-none-any.whl → 1.5.dev2448__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
monai/__init__.py CHANGED
@@ -136,4 +136,4 @@ except BaseException:
136
136
 
137
137
  if MONAIEnvVars.debug():
138
138
  raise
139
- __commit_id__ = "13b96aedc48ad2da16149490b06a1a6bd8361335"
139
+ __commit_id__ = "44e249d7d492d858199acfca1c948faa5aa33763"
monai/_version.py CHANGED
@@ -8,11 +8,11 @@ import json
8
8
 
9
9
  version_json = '''
10
10
  {
11
- "date": "2024-11-17T02:30:32+0000",
11
+ "date": "2024-12-01T02:35:43+0000",
12
12
  "dirty": false,
13
13
  "error": null,
14
- "full-revisionid": "218216250ce297265400abe56ee915898d75a2ec",
15
- "version": "1.5.dev2446"
14
+ "full-revisionid": "d4ff1455cf46b35e4dcfb6f57d54b0738b39f738",
15
+ "version": "1.5.dev2448"
16
16
  }
17
17
  ''' # END VERSION_JSON
18
18
 
monai/bundle/__init__.py CHANGED
@@ -43,4 +43,4 @@ from .utils import (
43
43
  MACRO_KEY,
44
44
  load_bundle_config,
45
45
  )
46
- from .workflows import BundleWorkflow, ConfigWorkflow
46
+ from .workflows import BundleWorkflow, ConfigWorkflow, PythonicWorkflow
@@ -192,6 +192,16 @@ class ReferenceResolver:
192
192
  """
193
193
  return self._resolve_one_item(id=id, **kwargs)
194
194
 
195
+ def remove_resolved_content(self, id: str) -> Any | None:
196
+ """
197
+ Remove the resolved ``ConfigItem`` by id.
198
+
199
+ Args:
200
+ id: id name of the expected item.
201
+
202
+ """
203
+ return self.resolved_content.pop(id) if id in self.resolved_content else None
204
+
195
205
  @classmethod
196
206
  def normalize_id(cls, id: str | int) -> str:
197
207
  """
monai/bundle/workflows.py CHANGED
@@ -44,12 +44,18 @@ class BundleWorkflow(ABC):
44
44
  workflow_type: specifies the workflow type: "train" or "training" for a training workflow,
45
45
  or "infer", "inference", "eval", "evaluation" for a inference workflow,
46
46
  other unsupported string will raise a ValueError.
47
- default to `train` for train workflow.
47
+ default to `None` for only using meta properties.
48
48
  workflow: specifies the workflow type: "train" or "training" for a training workflow,
49
49
  or "infer", "inference", "eval", "evaluation" for a inference workflow,
50
50
  other unsupported string will raise a ValueError.
51
51
  default to `None` for common workflow.
52
- properties_path: the path to the JSON file of properties.
52
+ properties_path: the path to the JSON file of properties. If `workflow_type` is specified, properties will be
53
+ loaded from the file based on the provided `workflow_type` and meta. If no `workflow_type` is specified,
54
+ properties will default to loading from "meta". If `properties_path` is None, default properties
55
+ will be sourced from "monai/bundle/properties.py" based on the workflow_type:
56
+ For a training workflow, properties load from `TrainProperties` and `MetaProperties`.
57
+ For a inference workflow, properties load from `InferProperties` and `MetaProperties`.
58
+ For workflow_type = None : only `MetaProperties` will be loaded.
53
59
  meta_file: filepath of the metadata file, if this is a list of file paths, their contents will be merged in order.
54
60
  logging_file: config file for `logging` module in the program. for more details:
55
61
  https://docs.python.org/3/library/logging.config.html#logging.config.fileConfig.
@@ -97,29 +103,50 @@ class BundleWorkflow(ABC):
97
103
  meta_file = None
98
104
 
99
105
  workflow_type = workflow if workflow is not None else workflow_type
100
- if workflow_type is None and properties_path is None:
101
- self.properties = copy(MetaProperties)
102
- self.workflow_type = None
103
- self.meta_file = meta_file
104
- return
106
+ if workflow_type is not None:
107
+ if workflow_type.lower() in self.supported_train_type:
108
+ workflow_type = "train"
109
+ elif workflow_type.lower() in self.supported_infer_type:
110
+ workflow_type = "infer"
111
+ else:
112
+ raise ValueError(f"Unsupported workflow type: '{workflow_type}'.")
113
+
105
114
  if properties_path is not None:
106
115
  properties_path = Path(properties_path)
107
116
  if not properties_path.is_file():
108
117
  raise ValueError(f"Property file {properties_path} does not exist.")
109
118
  with open(properties_path) as json_file:
110
- self.properties = json.load(json_file)
111
- self.workflow_type = None
112
- self.meta_file = meta_file
113
- return
114
- if workflow_type.lower() in self.supported_train_type: # type: ignore[union-attr]
115
- self.properties = {**TrainProperties, **MetaProperties}
116
- self.workflow_type = "train"
117
- elif workflow_type.lower() in self.supported_infer_type: # type: ignore[union-attr]
118
- self.properties = {**InferProperties, **MetaProperties}
119
- self.workflow_type = "infer"
119
+ try:
120
+ properties = json.load(json_file)
121
+ self.properties: dict = {}
122
+ if workflow_type is not None and workflow_type in properties:
123
+ self.properties = properties[workflow_type]
124
+ if "meta" in properties:
125
+ self.properties.update(properties["meta"])
126
+ elif workflow_type is None:
127
+ if "meta" in properties:
128
+ self.properties = properties["meta"]
129
+ logger.info(
130
+ "No workflow type specified, default to load meta properties from property file."
131
+ )
132
+ else:
133
+ logger.warning("No 'meta' key found in properties while workflow_type is None.")
134
+ except KeyError as e:
135
+ raise ValueError(f"{workflow_type} not found in property file {properties_path}") from e
136
+ except json.JSONDecodeError as e:
137
+ raise ValueError(f"Error decoding JSON from property file {properties_path}") from e
120
138
  else:
121
- raise ValueError(f"Unsupported workflow type: '{workflow_type}'.")
139
+ if workflow_type == "train":
140
+ self.properties = {**TrainProperties, **MetaProperties}
141
+ elif workflow_type == "infer":
142
+ self.properties = {**InferProperties, **MetaProperties}
143
+ elif workflow_type is None:
144
+ self.properties = copy(MetaProperties)
145
+ logger.info("No workflow type and property file specified, default to 'meta' properties.")
146
+ else:
147
+ raise ValueError(f"Unsupported workflow type: '{workflow_type}'.")
122
148
 
149
+ self.workflow_type = workflow_type
123
150
  self.meta_file = meta_file
124
151
 
125
152
  @abstractmethod
@@ -226,6 +253,124 @@ class BundleWorkflow(ABC):
226
253
  return [n for n, p in self.properties.items() if p.get(BundleProperty.REQUIRED, False) and not hasattr(self, n)]
227
254
 
228
255
 
256
+ class PythonicWorkflow(BundleWorkflow):
257
+ """
258
+ Base class for the pythonic workflow specification in bundle, it can be a training, evaluation or inference workflow.
259
+ It defines the basic interfaces for the bundle workflow behavior: `initialize`, `finalize`, etc.
260
+ This also provides the interface to get / set public properties to interact with a bundle workflow through
261
+ defined `get_<property>` accessor methods or directly defining members of the object.
262
+ For how to set the properties, users can define the `_set_<property>` methods or directly set the members of the object.
263
+ The `initialize` method is called to set up the workflow before running. This method sets up internal state
264
+ and prepares properties. If properties are modified after the workflow has been initialized, `self._is_initialized`
265
+ is set to `False`. Before running the workflow again, `initialize` should be called to ensure that the workflow is
266
+ properly set up with the new property values.
267
+
268
+ Args:
269
+ workflow_type: specifies the workflow type: "train" or "training" for a training workflow,
270
+ or "infer", "inference", "eval", "evaluation" for a inference workflow,
271
+ other unsupported string will raise a ValueError.
272
+ default to `None` for only using meta properties.
273
+ workflow: specifies the workflow type: "train" or "training" for a training workflow,
274
+ or "infer", "inference", "eval", "evaluation" for a inference workflow,
275
+ other unsupported string will raise a ValueError.
276
+ default to `None` for common workflow.
277
+ properties_path: the path to the JSON file of properties. If `workflow_type` is specified, properties will be
278
+ loaded from the file based on the provided `workflow_type` and meta. If no `workflow_type` is specified,
279
+ properties will default to loading from "meta". If `properties_path` is None, default properties
280
+ will be sourced from "monai/bundle/properties.py" based on the workflow_type:
281
+ For a training workflow, properties load from `TrainProperties` and `MetaProperties`.
282
+ For a inference workflow, properties load from `InferProperties` and `MetaProperties`.
283
+ For workflow_type = None : only `MetaProperties` will be loaded.
284
+ config_file: path to the config file, typically used to store hyperparameters.
285
+ meta_file: filepath of the metadata file, if this is a list of file paths, their contents will be merged in order.
286
+ logging_file: config file for `logging` module in the program. for more details:
287
+ https://docs.python.org/3/library/logging.config.html#logging.config.fileConfig.
288
+
289
+ """
290
+
291
+ supported_train_type: tuple = ("train", "training")
292
+ supported_infer_type: tuple = ("infer", "inference", "eval", "evaluation")
293
+
294
+ def __init__(
295
+ self,
296
+ workflow_type: str | None = None,
297
+ properties_path: PathLike | None = None,
298
+ config_file: str | Sequence[str] | None = None,
299
+ meta_file: str | Sequence[str] | None = None,
300
+ logging_file: str | None = None,
301
+ **override: Any,
302
+ ):
303
+ meta_file = str(Path(os.getcwd()) / "metadata.json") if meta_file is None else meta_file
304
+ super().__init__(
305
+ workflow_type=workflow_type, properties_path=properties_path, meta_file=meta_file, logging_file=logging_file
306
+ )
307
+ self._props_vals: dict = {}
308
+ self._set_props_vals: dict = {}
309
+ self.parser = ConfigParser()
310
+ if config_file is not None:
311
+ self.parser.read_config(f=config_file)
312
+ if self.meta_file is not None:
313
+ self.parser.read_meta(f=self.meta_file)
314
+
315
+ # the rest key-values in the _args are to override config content
316
+ self.parser.update(pairs=override)
317
+ self._is_initialized: bool = False
318
+
319
+ def initialize(self, *args: Any, **kwargs: Any) -> Any:
320
+ """
321
+ Initialize the bundle workflow before running.
322
+ """
323
+ self._props_vals = {}
324
+ self._is_initialized = True
325
+
326
+ def _get_property(self, name: str, property: dict) -> Any:
327
+ """
328
+ With specified property name and information, get the expected property value.
329
+ If the property is already generated, return from the bucket directly.
330
+ If user explicitly set the property, return it directly.
331
+ Otherwise, generate the expected property as a class private property with prefix "_".
332
+
333
+ Args:
334
+ name: the name of target property.
335
+ property: other information for the target property, defined in `TrainProperties` or `InferProperties`.
336
+ """
337
+ if not self._is_initialized:
338
+ raise RuntimeError("Please execute 'initialize' before getting any properties.")
339
+ value = None
340
+ if name in self._set_props_vals:
341
+ value = self._set_props_vals[name]
342
+ elif name in self._props_vals:
343
+ value = self._props_vals[name]
344
+ elif name in self.parser.config[self.parser.meta_key]: # type: ignore[index]
345
+ id = self.properties.get(name, None).get(BundlePropertyConfig.ID, None)
346
+ value = self.parser[id]
347
+ else:
348
+ try:
349
+ value = getattr(self, f"get_{name}")()
350
+ except AttributeError as e:
351
+ if property[BundleProperty.REQUIRED]:
352
+ raise ValueError(
353
+ f"unsupported property '{name}' is required in the bundle properties,"
354
+ f"need to implement a method 'get_{name}' to provide the property."
355
+ ) from e
356
+ self._props_vals[name] = value
357
+ return value
358
+
359
+ def _set_property(self, name: str, property: dict, value: Any) -> Any:
360
+ """
361
+ With specified property name and information, set value for the expected property.
362
+ Stores user-reset initialized objects that should not be re-initialized and marks the workflow as not initialized.
363
+
364
+ Args:
365
+ name: the name of target property.
366
+ property: other information for the target property, defined in `TrainProperties` or `InferProperties`.
367
+ value: value to set for the property.
368
+
369
+ """
370
+ self._set_props_vals[name] = value
371
+ self._is_initialized = False
372
+
373
+
229
374
  class ConfigWorkflow(BundleWorkflow):
230
375
  """
231
376
  Specification for the config-based bundle workflow.
@@ -262,7 +407,13 @@ class ConfigWorkflow(BundleWorkflow):
262
407
  or "infer", "inference", "eval", "evaluation" for a inference workflow,
263
408
  other unsupported string will raise a ValueError.
264
409
  default to `None` for common workflow.
265
- properties_path: the path to the JSON file of properties.
410
+ properties_path: the path to the JSON file of properties. If `workflow_type` is specified, properties will be
411
+ loaded from the file based on the provided `workflow_type` and meta. If no `workflow_type` is specified,
412
+ properties will default to loading from "train". If `properties_path` is None, default properties
413
+ will be sourced from "monai/bundle/properties.py" based on the workflow_type:
414
+ For a training workflow, properties load from `TrainProperties` and `MetaProperties`.
415
+ For a inference workflow, properties load from `InferProperties` and `MetaProperties`.
416
+ For workflow_type = None : only `MetaProperties` will be loaded.
266
417
  override: id-value pairs to override or add the corresponding config content.
267
418
  e.g. ``--net#input_chns 42``, ``--net %/data/other.json#net_arg``
268
419
 
@@ -324,7 +475,6 @@ class ConfigWorkflow(BundleWorkflow):
324
475
  self.parser.read_config(f=config_file)
325
476
  if self.meta_file is not None:
326
477
  self.parser.read_meta(f=self.meta_file)
327
-
328
478
  # the rest key-values in the _args are to override config content
329
479
  self.parser.update(pairs=override)
330
480
  self.init_id = init_id
@@ -394,8 +544,23 @@ class ConfigWorkflow(BundleWorkflow):
394
544
  ret.extend(wrong_props)
395
545
  return ret
396
546
 
397
- def _run_expr(self, id: str, **kwargs: dict) -> Any:
398
- return self.parser.get_parsed_content(id, **kwargs) if id in self.parser else None
547
+ def _run_expr(self, id: str, **kwargs: dict) -> list[Any]:
548
+ """
549
+ Evaluate the expression or expression list given by `id`. The resolved values from the evaluations are not stored,
550
+ allowing this to be evaluated repeatedly (eg. in streaming applications) without restarting the hosting process.
551
+ """
552
+ ret = []
553
+ if id in self.parser:
554
+ # suppose all the expressions are in a list, run and reset the expressions
555
+ if isinstance(self.parser[id], list):
556
+ for i in range(len(self.parser[id])):
557
+ sub_id = f"{id}{ID_SEP_KEY}{i}"
558
+ ret.append(self.parser.get_parsed_content(sub_id, **kwargs))
559
+ self.parser.ref_resolver.remove_resolved_content(sub_id)
560
+ else:
561
+ ret.append(self.parser.get_parsed_content(id, **kwargs))
562
+ self.parser.ref_resolver.remove_resolved_content(id)
563
+ return ret
399
564
 
400
565
  def _get_prop_id(self, name: str, property: dict) -> Any:
401
566
  prop_id = property[BundlePropertyConfig.ID]
@@ -11,7 +11,7 @@
11
11
 
12
12
  from __future__ import annotations
13
13
 
14
- from typing import Tuple, Union
14
+ from typing import Optional, Tuple, Union
15
15
 
16
16
  import torch
17
17
  import torch.nn as nn
@@ -154,10 +154,12 @@ class SABlock(nn.Module):
154
154
  )
155
155
  self.input_size = input_size
156
156
 
157
- def forward(self, x):
157
+ def forward(self, x, attn_mask: Optional[torch.Tensor] = None):
158
158
  """
159
159
  Args:
160
160
  x (torch.Tensor): input tensor. B x (s_dim_1 * ... * s_dim_n) x C
161
+ attn_mask (torch.Tensor, optional): mask to apply to the attention matrix.
162
+ B x (s_dim_1 * ... * s_dim_n). Defaults to None.
161
163
 
162
164
  Return:
163
165
  torch.Tensor: B x (s_dim_1 * ... * s_dim_n) x C
@@ -176,7 +178,13 @@ class SABlock(nn.Module):
176
178
 
177
179
  if self.use_flash_attention:
178
180
  x = F.scaled_dot_product_attention(
179
- query=q, key=k, value=v, scale=self.scale, dropout_p=self.dropout_rate, is_causal=self.causal
181
+ query=q,
182
+ key=k,
183
+ value=v,
184
+ attn_mask=attn_mask,
185
+ scale=self.scale,
186
+ dropout_p=self.dropout_rate,
187
+ is_causal=self.causal,
180
188
  )
181
189
  else:
182
190
  att_mat = torch.einsum("blxd,blyd->blxy", q, k) * self.scale
@@ -186,10 +194,16 @@ class SABlock(nn.Module):
186
194
  att_mat = self.rel_positional_embedding(x, att_mat, q)
187
195
 
188
196
  if self.causal:
197
+ if attn_mask is not None:
198
+ raise ValueError("Causal attention does not support attention masks.")
189
199
  att_mat = att_mat.masked_fill(self.causal_mask[:, :, : x.shape[-2], : x.shape[-2]] == 0, float("-inf"))
190
200
 
191
- att_mat = att_mat.softmax(dim=-1)
201
+ if attn_mask is not None:
202
+ attn_mask = attn_mask.unsqueeze(1).unsqueeze(2)
203
+ attn_mask = attn_mask.expand(-1, self.num_heads, -1, -1)
204
+ att_mat = att_mat.masked_fill(attn_mask == 0, float("-inf"))
192
205
 
206
+ att_mat = att_mat.softmax(dim=-1)
193
207
  if self.save_attn:
194
208
  # no gradients and new tensor;
195
209
  # https://pytorch.org/docs/stable/generated/torch.Tensor.detach.html
@@ -90,8 +90,10 @@ class TransformerBlock(nn.Module):
90
90
  use_flash_attention=use_flash_attention,
91
91
  )
92
92
 
93
- def forward(self, x: torch.Tensor, context: Optional[torch.Tensor] = None) -> torch.Tensor:
94
- x = x + self.attn(self.norm1(x))
93
+ def forward(
94
+ self, x: torch.Tensor, context: Optional[torch.Tensor] = None, attn_mask: Optional[torch.Tensor] = None
95
+ ) -> torch.Tensor:
96
+ x = x + self.attn(self.norm1(x), attn_mask=attn_mask)
95
97
  if self.with_cross_attention:
96
98
  x = x + self.cross_attn(self.norm_cross_attn(x), context=context)
97
99
  x = x + self.mlp(self.norm2(x))
@@ -53,6 +53,7 @@ from .fullyconnectednet import FullyConnectedNet, VarFullyConnectedNet
53
53
  from .generator import Generator
54
54
  from .highresnet import HighResBlock, HighResNet
55
55
  from .hovernet import Hovernet, HoVernet, HoVerNet, HoverNet
56
+ from .masked_autoencoder_vit import MaskedAutoEncoderViT
56
57
  from .mednext import (
57
58
  MedNeXt,
58
59
  MedNext,
@@ -0,0 +1,211 @@
1
+ # Copyright (c) MONAI Consortium
2
+ # Licensed under the Apache License, Version 2.0 (the "License");
3
+ # you may not use this file except in compliance with the License.
4
+ # You may obtain a copy of the License at
5
+ # http://www.apache.org/licenses/LICENSE-2.0
6
+ # Unless required by applicable law or agreed to in writing, software
7
+ # distributed under the License is distributed on an "AS IS" BASIS,
8
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
9
+ # See the License for the specific language governing permissions and
10
+ # limitations under the License.
11
+
12
+ from __future__ import annotations
13
+
14
+ from collections.abc import Sequence
15
+
16
+ import numpy as np
17
+ import torch
18
+ import torch.nn as nn
19
+
20
+ from monai.networks.blocks.patchembedding import PatchEmbeddingBlock
21
+ from monai.networks.blocks.pos_embed_utils import build_sincos_position_embedding
22
+ from monai.networks.blocks.transformerblock import TransformerBlock
23
+ from monai.networks.layers import trunc_normal_
24
+ from monai.utils import ensure_tuple_rep
25
+ from monai.utils.module import look_up_option
26
+
27
+ SUPPORTED_POS_EMBEDDING_TYPES = {"none", "learnable", "sincos"}
28
+
29
+ __all__ = ["MaskedAutoEncoderViT"]
30
+
31
+
32
+ class MaskedAutoEncoderViT(nn.Module):
33
+ """
34
+ Masked Autoencoder (ViT), based on: "Kaiming et al.,
35
+ Masked Autoencoders Are Scalable Vision Learners <https://arxiv.org/abs/2111.06377>"
36
+ Only a subset of the patches passes through the encoder. The decoder tries to reconstruct
37
+ the masked patches, resulting in improved training speed.
38
+ """
39
+
40
+ def __init__(
41
+ self,
42
+ in_channels: int,
43
+ img_size: Sequence[int] | int,
44
+ patch_size: Sequence[int] | int,
45
+ hidden_size: int = 768,
46
+ mlp_dim: int = 512,
47
+ num_layers: int = 12,
48
+ num_heads: int = 12,
49
+ masking_ratio: float = 0.75,
50
+ decoder_hidden_size: int = 384,
51
+ decoder_mlp_dim: int = 512,
52
+ decoder_num_layers: int = 4,
53
+ decoder_num_heads: int = 12,
54
+ proj_type: str = "conv",
55
+ pos_embed_type: str = "sincos",
56
+ decoder_pos_embed_type: str = "sincos",
57
+ dropout_rate: float = 0.0,
58
+ spatial_dims: int = 3,
59
+ qkv_bias: bool = False,
60
+ save_attn: bool = False,
61
+ ) -> None:
62
+ """
63
+ Args:
64
+ in_channels: dimension of input channels or the number of channels for input.
65
+ img_size: dimension of input image.
66
+ patch_size: dimension of patch size
67
+ hidden_size: dimension of hidden layer. Defaults to 768.
68
+ mlp_dim: dimension of feedforward layer. Defaults to 512.
69
+ num_layers: number of transformer blocks. Defaults to 12.
70
+ num_heads: number of attention heads. Defaults to 12.
71
+ masking_ratio: ratio of patches to be masked. Defaults to 0.75.
72
+ decoder_hidden_size: dimension of hidden layer for decoder. Defaults to 384.
73
+ decoder_mlp_dim: dimension of feedforward layer for decoder. Defaults to 512.
74
+ decoder_num_layers: number of transformer blocks for decoder. Defaults to 4.
75
+ decoder_num_heads: number of attention heads for decoder. Defaults to 12.
76
+ proj_type: position embedding layer type. Defaults to "conv".
77
+ pos_embed_type: position embedding layer type. Defaults to "sincos".
78
+ decoder_pos_embed_type: position embedding layer type for decoder. Defaults to "sincos".
79
+ dropout_rate: fraction of the input units to drop. Defaults to 0.0.
80
+ spatial_dims: number of spatial dimensions. Defaults to 3.
81
+ qkv_bias: apply bias to the qkv linear layer in self attention block. Defaults to False.
82
+ save_attn: to make accessible the attention in self attention block. Defaults to False.
83
+ Examples::
84
+ # for single channel input with image size of (96,96,96), and sin-cos positional encoding
85
+ >>> net = MaskedAutoEncoderViT(in_channels=1, img_size=(96,96,96), patch_size=(16,16,16),
86
+ pos_embed_type='sincos')
87
+ # for 3-channel with image size of (128,128,128) and a learnable positional encoding
88
+ >>> net = MaskedAutoEncoderViT(in_channels=3, img_size=128, patch_size=16, pos_embed_type='learnable')
89
+ # for 3-channel with image size of (224,224) and a masking ratio of 0.25
90
+ >>> net = MaskedAutoEncoderViT(in_channels=3, img_size=(224,224), patch_size=(16,16), masking_ratio=0.25,
91
+ spatial_dims=2)
92
+ """
93
+
94
+ super().__init__()
95
+
96
+ if not (0 <= dropout_rate <= 1):
97
+ raise ValueError(f"dropout_rate should be between 0 and 1, got {dropout_rate}.")
98
+
99
+ if hidden_size % num_heads != 0:
100
+ raise ValueError("hidden_size should be divisible by num_heads.")
101
+
102
+ if decoder_hidden_size % decoder_num_heads != 0:
103
+ raise ValueError("decoder_hidden_size should be divisible by decoder_num_heads.")
104
+
105
+ self.patch_size = ensure_tuple_rep(patch_size, spatial_dims)
106
+ self.img_size = ensure_tuple_rep(img_size, spatial_dims)
107
+ self.spatial_dims = spatial_dims
108
+ for m, p in zip(self.img_size, self.patch_size):
109
+ if m % p != 0:
110
+ raise ValueError(f"patch_size={patch_size} should be divisible by img_size={img_size}.")
111
+
112
+ self.decoder_hidden_size = decoder_hidden_size
113
+
114
+ if masking_ratio <= 0 or masking_ratio >= 1:
115
+ raise ValueError(f"masking_ratio should be in the range (0, 1), got {masking_ratio}.")
116
+
117
+ self.masking_ratio = masking_ratio
118
+ self.cls_token = nn.Parameter(torch.zeros(1, 1, hidden_size))
119
+
120
+ self.patch_embedding = PatchEmbeddingBlock(
121
+ in_channels=in_channels,
122
+ img_size=img_size,
123
+ patch_size=patch_size,
124
+ hidden_size=hidden_size,
125
+ num_heads=num_heads,
126
+ proj_type=proj_type,
127
+ pos_embed_type=pos_embed_type,
128
+ dropout_rate=dropout_rate,
129
+ spatial_dims=self.spatial_dims,
130
+ )
131
+ blocks = [
132
+ TransformerBlock(hidden_size, mlp_dim, num_heads, dropout_rate, qkv_bias, save_attn)
133
+ for _ in range(num_layers)
134
+ ]
135
+ self.blocks = nn.Sequential(*blocks, nn.LayerNorm(hidden_size))
136
+
137
+ # decoder
138
+ self.decoder_embed = nn.Linear(hidden_size, decoder_hidden_size)
139
+
140
+ self.mask_tokens = nn.Parameter(torch.zeros(1, 1, decoder_hidden_size))
141
+
142
+ self.decoder_pos_embed_type = look_up_option(decoder_pos_embed_type, SUPPORTED_POS_EMBEDDING_TYPES)
143
+ self.decoder_pos_embedding = nn.Parameter(torch.zeros(1, self.patch_embedding.n_patches, decoder_hidden_size))
144
+
145
+ decoder_blocks = [
146
+ TransformerBlock(decoder_hidden_size, decoder_mlp_dim, decoder_num_heads, dropout_rate, qkv_bias, save_attn)
147
+ for _ in range(decoder_num_layers)
148
+ ]
149
+ self.decoder_blocks = nn.Sequential(*decoder_blocks, nn.LayerNorm(decoder_hidden_size))
150
+ self.decoder_pred = nn.Linear(decoder_hidden_size, int(np.prod(self.patch_size)) * in_channels)
151
+
152
+ self._init_weights()
153
+
154
+ def _init_weights(self):
155
+ """
156
+ similar to monai/networks/blocks/patchembedding.py for the decoder positional encoding and for mask and
157
+ classification tokens
158
+ """
159
+ if self.decoder_pos_embed_type == "none":
160
+ pass
161
+ elif self.decoder_pos_embed_type == "learnable":
162
+ trunc_normal_(self.decoder_pos_embedding, mean=0.0, std=0.02, a=-2.0, b=2.0)
163
+ elif self.decoder_pos_embed_type == "sincos":
164
+ grid_size = []
165
+ for in_size, pa_size in zip(self.img_size, self.patch_size):
166
+ grid_size.append(in_size // pa_size)
167
+
168
+ self.decoder_pos_embedding = build_sincos_position_embedding(
169
+ grid_size, self.decoder_hidden_size, self.spatial_dims
170
+ )
171
+
172
+ else:
173
+ raise ValueError(f"decoder_pos_embed_type {self.decoder_pos_embed_type} not supported.")
174
+
175
+ # initialize patch_embedding like nn.Linear (instead of nn.Conv2d)
176
+ trunc_normal_(self.mask_tokens, mean=0.0, std=0.02, a=-2.0, b=2.0)
177
+ trunc_normal_(self.cls_token, mean=0.0, std=0.02, a=-2.0, b=2.0)
178
+
179
+ def _masking(self, x, masking_ratio: float | None = None):
180
+ batch_size, num_tokens, _ = x.shape
181
+ percentage_to_keep = 1 - masking_ratio if masking_ratio is not None else 1 - self.masking_ratio
182
+ selected_indices = torch.multinomial(
183
+ torch.ones(batch_size, num_tokens), int(percentage_to_keep * num_tokens), replacement=False
184
+ )
185
+ x_masked = x[torch.arange(batch_size).unsqueeze(1), selected_indices] # gather the selected tokens
186
+ mask = torch.ones(batch_size, num_tokens, dtype=torch.int).to(x.device)
187
+ mask[torch.arange(batch_size).unsqueeze(-1), selected_indices] = 0
188
+
189
+ return x_masked, selected_indices, mask
190
+
191
+ def forward(self, x, masking_ratio: float | None = None):
192
+ x = self.patch_embedding(x)
193
+ x, selected_indices, mask = self._masking(x, masking_ratio=masking_ratio)
194
+
195
+ cls_tokens = self.cls_token.expand(x.shape[0], -1, -1)
196
+ x = torch.cat((cls_tokens, x), dim=1)
197
+
198
+ x = self.blocks(x)
199
+
200
+ # decoder
201
+ x = self.decoder_embed(x)
202
+
203
+ x_ = self.mask_tokens.repeat(x.shape[0], mask.shape[1], 1)
204
+ x_[torch.arange(x.shape[0]).unsqueeze(-1), selected_indices] = x[:, 1:, :] # no cls token
205
+ x_ = x_ + self.decoder_pos_embedding
206
+ x = torch.cat([x[:, :1, :], x_], dim=1)
207
+ x = self.decoder_blocks(x)
208
+ x = self.decoder_pred(x)
209
+
210
+ x = x[:, 1:, :]
211
+ return x, mask