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,17 @@
import torch.nn as nn
from torchvision import models
# Base model
from architectures.Model import Model
class ShuffleNet(Model):
def get(self):
m = models.shufflenet_v2_x1_0(weights=models.ShuffleNet_V2_X1_0_Weights.DEFAULT)
num_ftrs = m.fc.in_features
m.fc = nn.Linear(num_ftrs, self.size)
return m