Wrapper that extracts intermediate layer features from VTPModel.
| 107 | # ============================================================================ |
| 108 | |
| 109 | class FeatureExtractor(nn.Module): |
| 110 | """Wrapper that extracts intermediate layer features from VTPModel.""" |
| 111 | |
| 112 | def __init__(self, model: VTPModel, n_last_blocks: int, autocast_dtype: torch.dtype): |
| 113 | super().__init__() |
| 114 | self.model = model |
| 115 | self.model.eval() |
| 116 | self.n_last_blocks = n_last_blocks |
| 117 | self.autocast_dtype = autocast_dtype |
| 118 | |
| 119 | def forward(self, images: torch.Tensor) -> List[Tuple[torch.Tensor, torch.Tensor]]: |
| 120 | """Extract intermediate layer features. |
| 121 | |
| 122 | Returns: |
| 123 | List of (patch_tokens, cls_token) tuples for each requested layer |
| 124 | """ |
| 125 | with torch.inference_mode(): |
| 126 | with torch.amp.autocast(device_type='cuda', dtype=self.autocast_dtype): |
| 127 | features = self.model.get_intermediate_layers_feature( |
| 128 | images, n=self.n_last_blocks, return_class_token=True |
| 129 | ) |
| 130 | return features |
| 131 | |
| 132 | |
| 133 | # ============================================================================ |
no outgoing calls