36 lines
1.1 KiB
Python
36 lines
1.1 KiB
Python
|
|
import torch.nn as nn
|
|
from torchvision import models
|
|
|
|
# Base model
|
|
from architectures.Model import Model
|
|
|
|
class ResNet50(Model):
|
|
# NOTE:
|
|
# This model had it's best performance with the following configs
|
|
# numbre of classes
|
|
# CLASS_SIZE = 20
|
|
# BATCH_SIZE = 16
|
|
# SAMPLE_SIZE = 30
|
|
# TRAINING_SMPLE = 28
|
|
# LR_RATE = 0.0001
|
|
# EPOCHS = 15
|
|
# RESOLUTION = 224
|
|
# NOTE: But it may be a one time thing.
|
|
# because testing again didn't repeat
|
|
|
|
def get(self):
|
|
m = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
|
|
|
|
# freez all layers
|
|
#for param in m.parameters():
|
|
#param.requires_grad = False
|
|
|
|
# unfreez the last two
|
|
# NOTE: Freezing everything and unfrizing the last 3 yeilded the best performance
|
|
#for param in m.layer2.parameters(): param.requires_grad = True
|
|
#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 |