File size: 1,043 Bytes
c849cf3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 | import torch
from transformers import PreTrainedModel
from .configuration_xmag import XMagConfig
class XMagModel(PreTrainedModel):
config_class = XMagConfig
base_model_prefix = "model"
def __init__(self, config):
super().__init__(config)
self.model = torch.hub.load("facebookresearch/dinov2", config.backbone_name)
def forward(self, pixel_values):
if pixel_values.ndim != 4:
raise ValueError("pixel_values must have shape (B, 3, H, W).")
_, channels, height, width = pixel_values.shape
if channels != 3:
raise ValueError("pixel_values must have 3 RGB channels.")
if height != self.config.image_size or width != self.config.image_size:
raise ValueError(f"pixel_values must be {self.config.image_size}x{self.config.image_size}.")
outputs = self.model.forward_features(pixel_values)
return {
"cls_embedding": outputs["x_norm_clstoken"],
"patch_embeddings": outputs["x_norm_patchtokens"],
}
|