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