15 lines
372 B
Python
15 lines
372 B
Python
|
|
|
|
import torch.nn as nn
|
|
from torchvision import models
|
|
|
|
# Base model
|
|
from architectures.Model import Model
|
|
|
|
class WideResNet(Model):
|
|
|
|
def get(self):
|
|
# wide_resnet50_2 is a common high-performance choice
|
|
m = models.wide_resnet50_2(weights=models.Wide_ResNet50_2_Weights.DEFAULT)
|
|
m.fc = nn.Linear(m.fc.in_features, self.size)
|
|
return m |