94 lines
3.1 KiB
Python
94 lines
3.1 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import os
|
|
import unittest
|
|
|
|
import torch
|
|
|
|
from .data.model_library import (DefaultModelLibrary, MMClsModelLibrary,
|
|
MMDetModelLibrary, MMModelLibrary,
|
|
MMPoseModelLibrary, MMSegModelLibrary,
|
|
ModelGenerator, TorchModelLibrary)
|
|
from .data.models import SingleLineModel
|
|
from .data.tracer_passed_models import (BackwardPassedModelManager,
|
|
FxPassedModelManager)
|
|
|
|
TEST_DATA = os.getenv('TEST_DATA') == 'true'
|
|
|
|
|
|
class TestModelLibrary(unittest.TestCase):
|
|
|
|
def test_mmcls(self):
|
|
if not TEST_DATA:
|
|
self.skipTest('not test data to save time.')
|
|
library = MMClsModelLibrary(exclude=['cutmax', 'cifar'])
|
|
self.assertTrue(library.is_default_includes_cover_all_models())
|
|
|
|
def test_defaul_library(self):
|
|
if not TEST_DATA:
|
|
self.skipTest('not test data to save time.')
|
|
library = DefaultModelLibrary()
|
|
self.assertTrue(library.is_default_includes_cover_all_models())
|
|
|
|
def test_torchlibrary(self):
|
|
if not TEST_DATA:
|
|
self.skipTest('not test data to save time.')
|
|
library = TorchModelLibrary()
|
|
self.assertTrue(library.is_default_includes_cover_all_models())
|
|
|
|
def test_mmdet(self):
|
|
if not TEST_DATA:
|
|
self.skipTest('not test data to save time.')
|
|
library = MMDetModelLibrary()
|
|
self.assertTrue(library.is_default_includes_cover_all_models())
|
|
|
|
def test_mmseg(self):
|
|
if not TEST_DATA:
|
|
self.skipTest('not test data to save time.')
|
|
library = MMSegModelLibrary()
|
|
print(library.short_names())
|
|
|
|
self.assertTrue(library.is_default_includes_cover_all_models())
|
|
|
|
# New
|
|
def test_mmpose(self):
|
|
if not TEST_DATA:
|
|
self.skipTest('not test data to save time.')
|
|
library = MMPoseModelLibrary()
|
|
print(library.short_names())
|
|
self.assertTrue(library.is_default_includes_cover_all_models())
|
|
|
|
def test_get_model_by_config(self):
|
|
config = 'mmcls::resnet/resnet34_8xb32_in1k.py'
|
|
Model = MMModelLibrary.get_model_from_path(config)
|
|
_ = Model()
|
|
|
|
def test_passed_models(self):
|
|
try:
|
|
print(FxPassedModelManager().include_models())
|
|
print(BackwardPassedModelManager().include_models())
|
|
except Exception:
|
|
self.fail()
|
|
|
|
|
|
class TestModels(unittest.TestCase):
|
|
|
|
def _test_a_model(self, Model):
|
|
model = Model()
|
|
x = torch.rand(2, 3, 224, 224)
|
|
y = model(x)
|
|
self.assertSequenceEqual(y.shape, [2, 1000])
|
|
|
|
def test_models(self):
|
|
library = DefaultModelLibrary()
|
|
for Model in library.include_models():
|
|
with self.subTest(model=Model):
|
|
self._test_a_model(Model)
|
|
|
|
def test_generator(self):
|
|
Model = ModelGenerator('model', SingleLineModel)
|
|
model = Model()
|
|
model.eval()
|
|
self.assertEqual(model.training, False)
|
|
model.train()
|
|
self.assertEqual(model.training, True)
|