17 lines
354 B
Python
17 lines
354 B
Python
|
|
|
|
|
|
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 |