Spaces:
Sleeping
Sleeping
fix indent
Browse files
model.py
CHANGED
|
@@ -3,15 +3,15 @@ import torchvision
|
|
| 3 |
from torch import nn
|
| 4 |
|
| 5 |
def create_efficientnet(output_shape: int):
|
| 6 |
-
|
| 7 |
-
|
| 8 |
|
| 9 |
-
|
| 10 |
-
|
| 11 |
|
| 12 |
-
|
| 13 |
-
|
| 14 |
-
|
| 15 |
-
|
| 16 |
|
| 17 |
return model, weights.transforms()
|
|
|
|
| 3 |
from torch import nn
|
| 4 |
|
| 5 |
def create_efficientnet(output_shape: int):
|
| 6 |
+
weights = torchvision.models.EfficientNet_V2_L_Weights.IMAGENET1K_V1.DEFAULT
|
| 7 |
+
model = torchvision.models.efficientnet_v2_l(weights=weights)
|
| 8 |
|
| 9 |
+
for param in model.parameters():
|
| 10 |
+
param.requires_grad = False
|
| 11 |
|
| 12 |
+
model.classifier = nn.Sequential(
|
| 13 |
+
nn.Dropout(p=0.10, inplace=True),
|
| 14 |
+
nn.Linear(in_features=1280, out_features=output_shape)
|
| 15 |
+
)
|
| 16 |
|
| 17 |
return model, weights.transforms()
|