31 lines
930 B
Python
31 lines
930 B
Python
|
|
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 |