Add NaViT: Native Resolution Vision Transformer with Patch n' Pack#9011
Add NaViT: Native Resolution Vision Transformer with Patch n' Pack#9011vikashg wants to merge 6 commits into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (6)
🚧 Files skipped from review as they are similar to previous changes (5)
📝 WalkthroughWalkthroughAdds the NaViT native-resolution vision transformer for variable-resolution 2D and 3D inputs. The implementation supports patch extraction, token dropout, sequence grouping, packed attention masks, factorized positional embeddings, transformer blocks, attention pooling, and classification. NaViT is exported from Estimated code review effort: 4 (Complex) | ~60 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://cdn.github.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
monai/networks/nets/navit.py (1)
427-446: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
forward()docstring omits theRaisessection.
forward()raisesValueErrorfor wrong ndim (line 475), wrong channel count (line 480), and non-divisible spatial dims (line 484), but none of this is documented. As per path instructions, docstrings should describe raised exceptions in the appropriate Google-style section.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@monai/networks/nets/navit.py` around lines 427 - 446, The forward() docstring in NaViT should document its ValueError cases. Add a Google-style Raises section describing ValueError for invalid tensor dimensionality, mismatched channel count, or spatial dimensions not divisible by patch_size.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@monai/networks/nets/navit.py`:
- Around line 344-348: Validate each dimension of image_size in the constructor
before building pos_embed_axes, requiring exact divisibility by patch_size and
raising a clear error otherwise. Use image_size_t and preserve the existing
positional-embedding table sizing only for valid dimensions.
- Around line 314-321: Update the constructor validation near the
hidden_size/num_heads divisibility check to explicitly reject non-positive
num_heads with ValueError before evaluating hidden_size % num_heads. Preserve
the existing ValueError for non-divisible positive values and the documented
validation behavior.
- Around line 492-497: Gate the token-dropout block in the NaViT forward path on
both self.calc_token_dropout being set and self.training, so dropout is applied
only during training. Keep the existing dropout fraction, index selection, and
sequence/position filtering unchanged when training is active.
- Around line 206-218: Update the forward method to preserve the post-attention
residual in x, pass only its normalized value to self.mlp, and add the MLP
output back to the unnormalized residual. Keep the attention residual and
existing return shape unchanged.
---
Nitpick comments:
In `@monai/networks/nets/navit.py`:
- Around line 427-446: The forward() docstring in NaViT should document its
ValueError cases. Add a Google-style Raises section describing ValueError for
invalid tensor dimensionality, mismatched channel count, or spatial dimensions
not divisible by patch_size.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 91dc08cb-fbc3-4e40-8d78-11336fac95b2
📒 Files selected for processing (6)
CHANGELOG.mddocs/source/networks.rstdocs/source/whatsnew_1_5_2.mdmonai/networks/nets/__init__.pymonai/networks/nets/navit.pytests/networks/nets/test_navit.py
| def forward(self, x: Tensor, attn_mask: Tensor | None = None) -> Tensor: | ||
| """ | ||
| Args: | ||
| x: input tensor of shape ``(B, N, C)``. | ||
| attn_mask: packed-sequence attention mask of shape ``(B, 1, N, N)``. | ||
|
|
||
| Returns: | ||
| Tensor of shape ``(B, N, C)``. | ||
| """ | ||
| x = self.attn(x, attn_mask=attn_mask) + x | ||
| x = self.norm(x) | ||
| x = self.mlp(x) + x | ||
| return x |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Residual stream gets overwritten by the normed value before the MLP add.
x = self.norm(x) at line 216 reassigns x, so the line 217 residual add (self.mlp(x) + x) adds the MLP output to the normalized value, not the raw pre-MLP residual sum. Standard pre-norm blocks (and the class docstring's own description "pre-norm self-attention followed by pre-norm MLP") keep the raw residual untouched and only feed the normed copy into the sublayer, e.g. x = x + mlp(norm(x)). As written, every block permanently collapses the residual stream through LayerNorm, discarding the growing-magnitude residual info deeper networks rely on for stable pre-norm training.
🔧 Proposed fix
def forward(self, x: Tensor, attn_mask: Tensor | None = None) -> Tensor:
x = self.attn(x, attn_mask=attn_mask) + x
- x = self.norm(x)
- x = self.mlp(x) + x
- return x
+ x = self.mlp(self.norm(x)) + x
+ return x📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def forward(self, x: Tensor, attn_mask: Tensor | None = None) -> Tensor: | |
| """ | |
| Args: | |
| x: input tensor of shape ``(B, N, C)``. | |
| attn_mask: packed-sequence attention mask of shape ``(B, 1, N, N)``. | |
| Returns: | |
| Tensor of shape ``(B, N, C)``. | |
| """ | |
| x = self.attn(x, attn_mask=attn_mask) + x | |
| x = self.norm(x) | |
| x = self.mlp(x) + x | |
| return x | |
| def forward(self, x: Tensor, attn_mask: Tensor | None = None) -> Tensor: | |
| """ | |
| Args: | |
| x: input tensor of shape ``(B, N, C)``. | |
| attn_mask: packed-sequence attention mask of shape ``(B, 1, N, N)``. | |
| Returns: | |
| Tensor of shape ``(B, N, C)``. | |
| """ | |
| x = self.attn(x, attn_mask=attn_mask) + x | |
| x = self.mlp(self.norm(x)) + x | |
| return x |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@monai/networks/nets/navit.py` around lines 206 - 218, Update the forward
method to preserve the post-attention residual in x, pass only its normalized
value to self.mlp, and add the MLP output back to the unnormalized residual.
Keep the attention residual and existing return shape unchanged.
| if spatial_dims not in (2, 3): | ||
| raise ValueError("spatial_dims must be 2 or 3.") | ||
| if not (0 <= dropout_rate <= 1): | ||
| raise ValueError("dropout_rate should be between 0 and 1.") | ||
| if not (0 <= emb_dropout_rate <= 1): | ||
| raise ValueError("emb_dropout_rate should be between 0 and 1.") | ||
| if hidden_size % num_heads != 0: | ||
| raise ValueError("hidden_size should be divisible by num_heads.") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
num_heads=0 raises ZeroDivisionError, not ValueError.
Line 320 (hidden_size % num_heads) crashes with ZeroDivisionError when num_heads=0, contradicting the documented "Raises: ValueError" contract for invalid constructor args.
🔧 Proposed fix
- if hidden_size % num_heads != 0:
+ if num_heads <= 0:
+ raise ValueError("num_heads must be a positive integer.")
+ if hidden_size % num_heads != 0:
raise ValueError("hidden_size should be divisible by num_heads.")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if spatial_dims not in (2, 3): | |
| raise ValueError("spatial_dims must be 2 or 3.") | |
| if not (0 <= dropout_rate <= 1): | |
| raise ValueError("dropout_rate should be between 0 and 1.") | |
| if not (0 <= emb_dropout_rate <= 1): | |
| raise ValueError("emb_dropout_rate should be between 0 and 1.") | |
| if hidden_size % num_heads != 0: | |
| raise ValueError("hidden_size should be divisible by num_heads.") | |
| if spatial_dims not in (2, 3): | |
| raise ValueError("spatial_dims must be 2 or 3.") | |
| if not (0 <= dropout_rate <= 1): | |
| raise ValueError("dropout_rate should be between 0 and 1.") | |
| if not (0 <= emb_dropout_rate <= 1): | |
| raise ValueError("emb_dropout_rate should be between 0 and 1.") | |
| if num_heads <= 0: | |
| raise ValueError("num_heads must be a positive integer.") | |
| if hidden_size % num_heads != 0: | |
| raise ValueError("hidden_size should be divisible by num_heads.") |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@monai/networks/nets/navit.py` around lines 314 - 321, Update the constructor
validation near the hidden_size/num_heads divisibility check to explicitly
reject non-positive num_heads with ValueError before evaluating hidden_size %
num_heads. Preserve the existing ValueError for non-divisible positive values
and the documented validation behavior.
| # --- factorised positional embeddings (one table per spatial axis) --- | ||
| image_size_t = ensure_tuple_rep(image_size, spatial_dims) | ||
| self.pos_embed_axes = nn.ParameterList( | ||
| [nn.Parameter(torch.randn(img_d // patch_size, hidden_size)) for img_d in image_size_t] | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
No divisibility check for image_size against patch_size.
img_d // patch_size silently floors if image_size isn't divisible by patch_size, producing an undersized positional-embedding table with no warning. Actual images are validated in forward() (line 483), but the reference image_size used only for table sizing is not. Consider validating divisibility at construction time for a clearer failure mode.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@monai/networks/nets/navit.py` around lines 344 - 348, Validate each dimension
of image_size in the constructor before building pos_embed_axes, requiring exact
divisibility by patch_size and raising a clear error otherwise. Use image_size_t
and preserve the existing positional-embedding table sizing only for valid
dimensions.
| if self.calc_token_dropout is not None: | ||
| dropout_frac = self.calc_token_dropout(*spatial) | ||
| num_keep = max(1, int(seq.shape[0] * (1.0 - dropout_frac))) | ||
| keep_idx = torch.randn(seq.shape[0], device=device).topk(num_keep).indices | ||
| seq = seq[keep_idx] | ||
| pos = pos[keep_idx] |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Token dropout is applied even in eval mode, breaking deterministic inference.
self.calc_token_dropout is invoked unconditionally, so tokens are randomly dropped at inference time too, making net.eval() outputs non-deterministic across calls whenever token_dropout_prob is set. The reference NaViT implementation gates this explicitly: has_token_dropout = exists(self.calc_token_dropout) and self.training. test_token_dropout_callable (tests/networks/nets/test_navit.py lines 346-362) calls net.eval() but only checks output shape, so this gap isn't caught by the current suite.
🔧 Proposed fix
- if self.calc_token_dropout is not None:
+ if self.calc_token_dropout is not None and self.training:
dropout_frac = self.calc_token_dropout(*spatial)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if self.calc_token_dropout is not None: | |
| dropout_frac = self.calc_token_dropout(*spatial) | |
| num_keep = max(1, int(seq.shape[0] * (1.0 - dropout_frac))) | |
| keep_idx = torch.randn(seq.shape[0], device=device).topk(num_keep).indices | |
| seq = seq[keep_idx] | |
| pos = pos[keep_idx] | |
| if self.calc_token_dropout is not None and self.training: | |
| dropout_frac = self.calc_token_dropout(*spatial) | |
| num_keep = max(1, int(seq.shape[0] * (1.0 - dropout_frac))) | |
| keep_idx = torch.randn(seq.shape[0], device=device).topk(num_keep).indices | |
| seq = seq[keep_idx] | |
| pos = pos[keep_idx] |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@monai/networks/nets/navit.py` around lines 492 - 497, Gate the token-dropout
block in the NaViT forward path on both self.calc_token_dropout being set and
self.training, so dropout is applied only during training. Keep the existing
dropout fraction, index selection, and sequence/position filtering unchanged
when training is active.
Adds NaViT (monai.networks.nets.NaViT), a Vision Transformer that removes the fixed-resolution constraint of standard ViT by packing multiple variable-size images into a single sequence per batch element. Key features: - Patch n' Pack: multiple images concatenated into one sequence per group, with a per-image attention mask preventing cross-image attention - Factorised positional embeddings: separate learnable tables per spatial axis, allowing generalisation to unseen resolutions - Token dropout: configurable fraction of patch tokens dropped during training (float or callable) - Attention pooling: learned query attends over each image's tokens to produce a fixed-size per-image representation - QK normalisation: RMS normalisation on queries and keys (ViT-22B style) - 2D and 3D support: works for (C, H, W) and (C, H, W, D) inputs Changes: - monai/networks/nets/navit.py: new NaViT implementation - monai/networks/nets/__init__.py: export NaViT - tests/networks/nets/test_navit.py: 24 unit tests covering shape, variable resolutions, token dropout, auto-grouping, gradient flow, ill arguments, and forward validation - docs/source/networks.rst: autoclass entry - docs/source/whatsnew_1_5_2.md: feature description - CHANGELOG.md: entry under Unreleased Signed-off-by: Vikash Gupta <write2vikash@gmail.com>
for more information, see https://pre-commit.ci
Signed-off-by: Vikash Gupta <write2vikash@gmail.com>
NaViT depends on einops (optional dependency), so test_navit must be excluded from the minimal CI runner per CONTRIBUTING.md guidelines. Signed-off-by: Vikash Gupta <write2vikash@gmail.com>
…ling - Remove NaViT section from whatsnew_1_5_2.md: 1.5.2 was a security-only patch release; NaViT belongs in the upcoming release under [Unreleased] - CHANGELOG [Unreleased]: factorised -> factorized, normalisation -> normalization - navit.py docstrings and comments: normalisation -> normalization, generalisation -> generalization, factorised -> factorized (MONAI uses American English per CONTRIBUTING.md) Signed-off-by: Vikash Gupta <write2vikash@gmail.com>
Remediation for bot commit aa7a76a ([pre-commit.ci] auto fixes from pre-commit.com hooks) which removed an unused SABlock import from navit.py. That commit was made by pre-commit-ci[bot] and could not carry a Signed-off-by line. Signed-off-by: Vikash Gupta <write2vikash@gmail.com>
1 participant
Adds NaViT (monai.networks.nets.NaViT), a Vision Transformer that removes the fixed-resolution constraint of standard ViT by packing multiple variable-size images into a single sequence per batch element.
Key features:
Changes:
Fixes # .
Description
A few sentences describing the changes proposed in this pull request.
Types of changes
./runtests.sh -f -u --net --coverage../runtests.sh --quick --unittests --disttests.make htmlcommand in thedocs/folder.