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()
|