Upload directory

models/vit_kprpe/__init__.py

@@ -0,0 +1,65 @@
+from ..base import BaseModel
+from .vit import VisionTransformerWithKPRPE
+from torchvision import transforms
+
+
+class ViTKPRPEModel(BaseModel):
+
+
+ """
+ Vision Transformer for face recognition model with KeyPoint Relative Position Encoding (KP-RPE).
+
+ ```
+ @article{kim2024keypoint,
+ title={KeyPoint Relative Position Encoding for Face Recognition},
+ author={Kim, Minchul and Su, Yiyang and Liu, Feng and Jain, Anil and Liu, Xiaoming},
+ journal={CVPR},
+ year={2024}
+ }
+ ```
+ """
+ def __init__(self, net, config):
+ super(ViTKPRPEModel, self).__init__(config)
+ self.net = net
+
+
+ @classmethod
+ def from_config(cls, config):
+
+ if config.name == 'small':
+ net = VisionTransformerWithKPRPE(img_size=112, patch_size=8, num_classes=config.output_dim, embed_dim=512, depth=12,
+ mlp_ratio=5, num_heads=8, drop_path_rate=0.1, norm_layer="ln",
+ mask_ratio=config.mask_ratio, rpe_config=config.rpe_config)
+ elif config.name == 'base':
+ net = VisionTransformerWithKPRPE(img_size=112, patch_size=8, num_classes=config.output_dim, embed_dim=512, depth=24,
+ mlp_ratio=3, num_heads=16, drop_path_rate=0.1, norm_layer="ln",
+ mask_ratio=config.mask_ratio, rpe_config=config.rpe_config)
+ else:
+ raise NotImplementedError
+
+ model = cls(net, config)
+ model.eval()
+ return model
+
+ def forward(self, x, *args, **kwargs):
+ if self.input_color_flip:
+ x = x.flip(1)
+ return self.net(x, *args, **kwargs)
+
+ def make_train_transform(self):
+ transform = transforms.Compose([
+ transforms.ToTensor(),
+ transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+ ])
+ return transform
+
+ def make_test_transform(self):
+ transform = transforms.Compose([
+ transforms.ToTensor(),
+ transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+ ])
+ return transform
+
+def load_model(model_config):
+ model = ViTKPRPEModel.from_config(model_config)
+ return model