Spaces:
Sleeping
Sleeping
Philipp Normann commited on
Commit ·
90918bb
1
Parent(s): ff533bd
Remove unused model class
Browse files
model.py
DELETED
|
@@ -1,16 +0,0 @@
|
|
| 1 |
-
import lightning as L
|
| 2 |
-
import torch.nn as nn
|
| 3 |
-
from torchvision import models
|
| 4 |
-
|
| 5 |
-
|
| 6 |
-
class ScribbleItNet(L.LightningModule):
|
| 7 |
-
|
| 8 |
-
def __init__(self, num_classes):
|
| 9 |
-
super().__init__()
|
| 10 |
-
|
| 11 |
-
self.model = models.efficientnet_b0()
|
| 12 |
-
self.num_features = self.model.classifier[-1].in_features
|
| 13 |
-
self.model.classifier[-1] = nn.Linear(self.num_features, num_classes)
|
| 14 |
-
|
| 15 |
-
def forward(self, image):
|
| 16 |
-
return self.model(image)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|