Spaces:
Sleeping
Sleeping
File size: 844 Bytes
c08ab4e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
import torch
import torch.nn as nn
import torchvision.models as models
class ResNet18(nn.Module):
def __init__(self, num_classes=15):
super(ResNet18, self).__init__()
self.model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
# Replace the last fully connected layer
self.model.fc = nn.Linear(self.model.fc.in_features, num_classes)
def forward(self, x):
x = self.model.conv1(x)
x = self.model.bn1(x)
x = self.model.relu(x)
x = self.model.maxpool(x)
x = self.model.layer1(x)
x = self.model.layer2(x)
x = self.model.layer3(x)
x = self.model.layer4(x)
x = self.model.avgpool(x)
x = torch.flatten(x, 1) # <-- here torch.flatten is used, so torch must be imported
x = self.model.fc(x)
return x
|