15 lines
423 B
Python
15 lines
423 B
Python
|
|
import torch.nn as nn
|
|
from torchvision import models
|
|
from architectures.Model import Model
|
|
|
|
class DenseNet121(Model):
|
|
def get(self):
|
|
|
|
# load pretrained
|
|
m = models.densenet121(weights=models.DenseNet121_Weights.DEFAULT)
|
|
# will modify only the final layers
|
|
num_ftrs = m.classifier.in_features
|
|
m.classifier = nn.Linear(num_ftrs, self.size)
|
|
|
|
return m.to(self.device) |