Spaces:
Sleeping
Sleeping
import torch | |
import torch.nn as nn | |
N_EMOTIONS = 8 | |
N_CELEBRITIES = 17 | |
class CustomModel(nn.Module) : | |
def __init__(self,mode = 'emotion') : | |
super().__init__() | |
self.mode = mode | |
self.backbone = nn.Sequential( | |
#3x224x224 | |
nn.Conv2d(3, 64, kernel_size=3, stride=1, bias=False), | |
nn.BatchNorm2d(64), | |
nn.LeakyReLU(0.2, inplace=True), | |
# out: 64 x 222 x 222 | |
nn.Conv2d(64, 32, kernel_size=3, stride=1, bias=False), | |
nn.BatchNorm2d(32), | |
nn.LeakyReLU(0.2, inplace=True), | |
nn.MaxPool2d(kernel_size=2), | |
nn.Dropout(0.2), | |
# out: 32 x 110 x 110 | |
nn.Conv2d(32, 32, kernel_size=3, stride=1, bias=False), | |
nn.BatchNorm2d(32), | |
nn.LeakyReLU(0.2, inplace=True), | |
nn.MaxPool2d(kernel_size=2), | |
nn.Dropout(0.3), | |
# out: 32 x 54 x 54 | |
nn.Flatten(), | |
) | |
self.in_features = 32*54*54 | |
self.neck = nn.Sequential( | |
nn.Linear(self.in_features,128), | |
nn.ReLU(), | |
nn.Linear(128,64), | |
nn.ReLU() | |
) | |
self.emotion_classifier = nn.Linear(64,N_EMOTIONS) | |
self.celebrity_classifier = nn.Linear(64,N_CELEBRITIES) | |
def forward(self,image) : | |
features = self.backbone(image) | |
features = self.neck(features) | |
if self.mode=='emotion' : | |
emotion_logits = self.emotion_classifier(features) | |
return emotion_logits | |
elif self.mode=='celebrity' : | |
celebrity_logits = self.celebrity_classifier(features) | |
return celebrity_logits | |
else : | |
emotion_logits = self.emotion_classifier(features) | |
celebrity_logits = self.celebrity_classifier(features) | |
return emotion_logits,celebrity_logits | |
import torchvision.models as models | |
class ResNet50Model(nn.Module) : | |
def __init__(self,mode = 'emotion') : | |
super().__init__() | |
self.mode = mode | |
self.backbone = getattr(models, 'resnet50')(False) | |
self.in_features = 1000 | |
self.neck = nn.Sequential( | |
nn.Linear(self.in_features,128), | |
nn.ReLU(), | |
nn.Linear(128,64), | |
nn.ReLU() | |
) | |
self.emotion_classifier = nn.Linear(64,N_EMOTIONS) | |
self.celebrity_classifier = nn.Linear(64,N_CELEBRITIES) | |
def forward(self,image) : | |
features = self.backbone(image) | |
features = self.neck(features) | |
if self.mode=='emotion' : | |
emotion_logits = self.emotion_classifier(features) | |
return emotion_logits | |
elif self.mode=='celebrity' : | |
celebrity_logits = self.celebrity_classifier(features) | |
return celebrity_logits | |
else : | |
emotion_logits = self.emotion_classifier(features) | |
celebrity_logits = self.celebrity_classifier(features) | |
return emotion_logits,celebrity_logits | |