admin commited on
Commit
e9ebebe
·
1 Parent(s): 6d8ed08
Files changed (2) hide show
  1. model.py +7 -5
  2. requirements.txt +1 -1
model.py CHANGED
@@ -2,7 +2,8 @@ import torch
2
  import torch.nn as nn
3
  import torchvision.models as models
4
  from modelscope.msdatasets import MsDataset
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("不支持的模型")
21
 
22
  self.model = eval("models.%s()" % m_ver)
23
  linear_output = self._set_outsize()
@@ -38,9 +39,10 @@ class EvalNet:
38
  return backbone_list[0]
39
 
40
  def _model_info(self, m_ver: str):
41
- backbone_list = MsDataset.load(
42
- "monetjoe/cv_backbones",
43
- split="v1",
 
44
  )
45
  backbone = self._get_backbone(m_ver, backbone_list)
46
  m_type = str(backbone["type"])
 
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()
 
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"])
requirements.txt CHANGED
@@ -4,4 +4,4 @@ torchvision==0.21.0+cu118
4
  -f https://download.pytorch.org/whl/torchvision
5
  librosa
6
  matplotlib
7
- modelscope[framework]==1.24.0
 
4
  -f https://download.pytorch.org/whl/torchvision
5
  librosa
6
  matplotlib
7
+ modelscope[framework]==1.21.0