added more models for testing

This commit is contained in:
2026-05-05 21:04:33 +02:00
parent 1c04344ad6
commit 4cc9fa2bac
11 changed files with 143 additions and 21 deletions

View File

@@ -0,0 +1,19 @@
import torch.nn as nn
from torchvision import models
# Base model
from architectures.Model import Model
class EfficientNet(Model):
def get(self):
m = models.efficientnet_b1(weights=models.EfficientNet_B1_Weights.DEFAULT)
# Unfreeze the last block for a lighter touch
for param in m.features[-1].parameters(): param.requires_grad = True
# Standard classifier fix
m.classifier[1] = nn.Linear(m.classifier[1].in_features, self.size)
return m