Convert 3D patch tokens (B, N, C) to 4D format (B, C, H, W). Args: patch_tokens: Patch tokens with shape (B, N, C) patch_size: Size of each patch img_h: Original image height img_w: Original image width Returns: Reshaped tokens with shape (B, C,
(
patch_tokens: torch.Tensor,
patch_size: int,
img_h: int,
img_w: int
)
| 2 | |
| 3 | |
| 4 | def convert_patch_tokens_to_4d( |
| 5 | patch_tokens: torch.Tensor, |
| 6 | patch_size: int, |
| 7 | img_h: int, |
| 8 | img_w: int |
| 9 | ) -> torch.Tensor: |
| 10 | """ |
| 11 | Convert 3D patch tokens (B, N, C) to 4D format (B, C, H, W). |
| 12 | |
| 13 | Args: |
| 14 | patch_tokens: Patch tokens with shape (B, N, C) |
| 15 | patch_size: Size of each patch |
| 16 | img_h: Original image height |
| 17 | img_w: Original image width |
| 18 | |
| 19 | Returns: |
| 20 | Reshaped tokens with shape (B, C, H, W) |
| 21 | """ |
| 22 | B, N, C = patch_tokens.shape |
| 23 | feat_h = img_h // patch_size |
| 24 | feat_w = img_w // patch_size |
| 25 | |
| 26 | expected_patches = feat_h * feat_w |
| 27 | if N != expected_patches: |
| 28 | raise ValueError( |
| 29 | f"Patch tokens mismatch: got {N}, expected {expected_patches}" |
| 30 | ) |
| 31 | |
| 32 | return patch_tokens.transpose(1, 2).reshape(B, C, feat_h, feat_w) |
nothing calls this directly
no outgoing calls
no test coverage detected