admin commited on
Commit
6e9a888
·
1 Parent(s): b027a5c
Files changed (1) hide show
  1. model.py +9 -4
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 Exception("Unsupported model.")
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("Backbone name not found, using default option - alexnet.")
38
  return backbone_list[0]
39
 
40
  def _model_info(self, m_ver: str):
41
- backbone_list = load_dataset("monetjoe/cv_backbones", split="train")
 
 
 
 
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"])