Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -3,7 +3,7 @@ import torch
|
|
3 |
import torch.nn as nn
|
4 |
from torchvision import transforms, models
|
5 |
import pickle
|
6 |
-
|
7 |
|
8 |
with open('class_names.pkl', 'rb') as f:
|
9 |
class_names = pickle.load(f)
|
@@ -11,7 +11,7 @@ with open('class_names.pkl', 'rb') as f:
|
|
11 |
# 加载训练好的模型
|
12 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
13 |
|
14 |
-
model =
|
15 |
model.fc = nn.Sequential(
|
16 |
nn.Dropout(0.2),
|
17 |
nn.Linear(model.fc.in_features, len(class_names))
|
|
|
3 |
import torch.nn as nn
|
4 |
from torchvision import transforms, models
|
5 |
import pickle
|
6 |
+
from resnest.torch import resnest50
|
7 |
|
8 |
with open('class_names.pkl', 'rb') as f:
|
9 |
class_names = pickle.load(f)
|
|
|
11 |
# 加载训练好的模型
|
12 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
13 |
|
14 |
+
model = resnest50(pretrained=None)
|
15 |
model.fc = nn.Sequential(
|
16 |
nn.Dropout(0.2),
|
17 |
nn.Linear(model.fc.in_features, len(class_names))
|