Spaces:
Running
Running
admin
commited on
Commit
·
6e9a888
1
Parent(s):
b027a5c
upd model
Browse files
model.py
CHANGED
@@ -1,8 +1,9 @@
|
|
1 |
import torch
|
2 |
import torch.nn as nn
|
3 |
import torchvision.models as models
|
|
|
4 |
from datasets import load_dataset
|
5 |
-
from utils import MODEL_DIR
|
6 |
|
7 |
|
8 |
class EvalNet:
|
@@ -17,7 +18,7 @@ class EvalNet:
|
|
17 |
self.m_type, self.input_size = self._model_info(m_ver)
|
18 |
|
19 |
if not hasattr(models, m_ver):
|
20 |
-
raise
|
21 |
|
22 |
self.model = eval("models.%s()" % m_ver)
|
23 |
linear_output = self._set_outsize()
|
@@ -34,11 +35,15 @@ class EvalNet:
|
|
34 |
if ver == bb["ver"]:
|
35 |
return bb
|
36 |
|
37 |
-
print("
|
38 |
return backbone_list[0]
|
39 |
|
40 |
def _model_info(self, m_ver: str):
|
41 |
-
backbone_list =
|
|
|
|
|
|
|
|
|
42 |
backbone = self._get_backbone(m_ver, backbone_list)
|
43 |
m_type = str(backbone["type"])
|
44 |
input_size = int(backbone["input_size"])
|
|
|
1 |
import torch
|
2 |
import torch.nn as nn
|
3 |
import torchvision.models as models
|
4 |
+
from modelscope.msdatasets import MsDataset
|
5 |
from datasets import load_dataset
|
6 |
+
from utils import MODEL_DIR, EN_US
|
7 |
|
8 |
|
9 |
class EvalNet:
|
|
|
18 |
self.m_type, self.input_size = self._model_info(m_ver)
|
19 |
|
20 |
if not hasattr(models, m_ver):
|
21 |
+
raise ValueError("不支持的模型")
|
22 |
|
23 |
self.model = eval("models.%s()" % m_ver)
|
24 |
linear_output = self._set_outsize()
|
|
|
35 |
if ver == bb["ver"]:
|
36 |
return bb
|
37 |
|
38 |
+
print("未找到骨干网络名称,使用默认选项 - alexnet。")
|
39 |
return backbone_list[0]
|
40 |
|
41 |
def _model_info(self, m_ver: str):
|
42 |
+
backbone_list = (
|
43 |
+
load_dataset("monetjoe/cv_backbones", split="train")
|
44 |
+
if EN_US
|
45 |
+
else MsDataset.load("monetjoe/cv_backbones", split="v1")
|
46 |
+
)
|
47 |
backbone = self._get_backbone(m_ver, backbone_list)
|
48 |
m_type = str(backbone["type"])
|
49 |
input_size = int(backbone["input_size"])
|