File size: 2,153 Bytes
0f3fa05 | 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 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | import torch
import torch.nn as nn
import torch.nn.functional as F
import clip
class CLIP_ViT_B_32(nn.Module):
def __init__(self, pretrained=True):
super().__init__()
if pretrained:
self.model, _ = clip.load("ViT-B/32", device="cpu")
else:
self.model, _ = clip.load("ViT-B/32", device="cpu", jit=False)
# Remove the final classification layer
self.feature_dim = 512 # CLIP ViT-B/32 feature dimension
def freeze_weights(self):
"""Freeze all CLIP weights"""
for param in self.model.parameters():
param.requires_grad = False
def forward(self, x):
"""Extract features from CLIP"""
with torch.no_grad() if not self.training else torch.enable_grad():
features = self.model.encode_image(x)
return features.float()
class CosineClassifier(nn.Module):
def __init__(self, dim, num_classes, scale_factor=30):
super().__init__()
self.num_classes = num_classes
self.scale_factor = scale_factor
self.weights = nn.Parameter(torch.randn(num_classes, dim))
nn.init.xavier_uniform_(self.weights)
def forward(self, x):
# Normalize input features
x = F.normalize(x, dim=-1)
# Normalize weights
w = F.normalize(self.weights, dim=-1)
# Compute cosine similarity and scale
logits = torch.matmul(x, w.T) * self.scale_factor
return logits
class FaceFairnessModel(nn.Module):
def __init__(self, num_classes, scale_factor=30):
super().__init__()
# CLIP feature extractor
self.clip_model = CLIP_ViT_B_32(pretrained=True)
self.clip_model.freeze_weights()
# Cosine classifier
self.classifier = CosineClassifier(
dim=self.clip_model.feature_dim,
num_classes=num_classes,
scale_factor=scale_factor
)
def forward(self, x):
# Extract features with CLIP
features = self.clip_model(x)
# Classify with cosine classifier
logits = self.classifier(features)
return logits |