| from torchvision.models import resnet50, ResNet50_Weights |
| from transformers import PreTrainedModel |
| from .config import ResnetConfig |
| import torch.nn as nn |
|
|
| class ResNet50(nn.Module): |
| def __init__(self, ): |
| super().__init__() |
| self.cnn = resnet50(weights=ResNet50_Weights.IMAGENET1K_V1) |
| self.backbone = nn.Sequential(*list(self.cnn.children())[:-2]) |
| self.flaten = nn.Sequential(nn.AvgPool2d(kernel_size=7), nn.Flatten()) |
| self.fc_1 = nn.Linear(2048, 768) |
|
|
| def forward(self, x): |
| if len(x.shape) == 3: |
| x = x.unsqueeze(0) |
| x = self.backbone(x) |
| x = self.flaten(x) |
| x = self.fc_1(x) |
| x = x.squeeze(0) |
| return x |
| |
| class ResNet50AffectiveFeatureExtractor(PreTrainedModel): |
| config_class = ResnetConfig |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| self.model = ResNet50() |
| del self.model.cnn |
|
|
| def forward(self, tensor): |
| return self.model(tensor) |
