import torch.nn as nn
from torchvision import models

# Base model
from architectures.Model import Model

class GoogleNet(Model):
    
    def get(self):

        m = models.googlenet(weights=models.GoogLeNet_Weights.DEFAULT)

        # 1. Handle the two Auxiliary Classifiers
        # GoogLeNet has aux1 and aux2 to help training converge
        #if m.aux_logits:
            #m.aux1.fc = nn.Linear(m.aux1.fc.in_features, self.size)
            #m.aux2.fc = nn.Linear(m.aux2.fc.in_features, self.size)

        # 2. Handle the Main Classifier
        m.fc = nn.Linear(m.fc.in_features, self.size)
        
        #for param in m.parameters():
        #    param.requires_grad = False
            
        # Unfreezing the final stages for identity recognition
        #for name, param in m.named_parameters():
        #    if "inception5" in name or "fc" in name:
        #        param.requires_grad = True

        return m