added more models for testing
This commit is contained in:
14
architectures/GoogleNet.py
Normal file
14
architectures/GoogleNet.py
Normal file
@@ -0,0 +1,14 @@
|
||||
|
||||
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)
|
||||
m.fc = nn.Linear(m.fc.in_features, self.size)
|
||||
return m
|
||||
Reference in New Issue
Block a user