File size: 1,794 Bytes
ab687e7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
from pytorch_caney.models.build import build_model
from pytorch_caney.config import get_config
import unittest
import argparse
import logging
class TestBuildModel(unittest.TestCase):
def setUp(self):
# Initialize any required configuration here
config_path = 'pytorch_caney/' + \
'tests/config/test_config.yaml'
args = argparse.Namespace(cfg=config_path)
self.config = get_config(args)
self.logger = logging.getLogger("TestLogger")
self.logger.setLevel(logging.DEBUG)
def test_build_mim_model(self):
_ = build_model(self.config,
pretrain=True,
pretrain_method='mim',
logger=self.logger)
# Add assertions here to validate the returned 'model' instance
# For example: self.assertIsInstance(model, YourMimModelClass)
def test_build_swinv2_encoder(self):
_ = build_model(self.config, logger=self.logger)
# Add assertions here to validate the returned 'model' instance
# For example: self.assertIsInstance(model, SwinTransformerV2)
def test_build_unet_decoder(self):
self.config.defrost()
self.config.MODEL.DECODER = 'unet'
self.config.freeze()
_ = build_model(self.config, logger=self.logger)
# Add assertions here to validate the returned 'model' instance
# For example: self.assertIsInstance(model, YourUnetSwinModelClass)
def test_unknown_decoder_architecture(self):
self.config.defrost()
self.config.MODEL.DECODER = 'unknown_decoder'
self.config.freeze()
with self.assertRaises(NotImplementedError):
build_model(self.config, logger=self.logger)
if __name__ == '__main__':
unittest.main()
|