22 lines
585 B
Python
22 lines
585 B
Python
|
|
import torch.nn as nn
|
|
from torchvision import models
|
|
|
|
# Base model
|
|
from architectures.Model import Model
|
|
|
|
class ResNet18(Model):
|
|
|
|
def get(self):
|
|
m = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
|
|
|
|
# freez all layers
|
|
#for param in m.parameters():
|
|
# param.requires_grad = False
|
|
|
|
# unfreez the last two
|
|
#for param in m.layer3.parameters(): param.requires_grad = True
|
|
#for param in m.layer4.parameters(): param.requires_grad = True
|
|
|
|
m.fc = nn.Linear(m.fc.in_features, self.size)
|
|
return m |