19 lines
501 B
Python
19 lines
501 B
Python
|
|
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 |