File size: 2,488 Bytes
3133fdb |
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 52 53 54 55 56 57 |
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
import unittest
import hydra
from hydra.experimental import compose, initialize_config_module
from pytorchvideo_trainer.module.byol import BYOLModule
from pytorchvideo_trainer.module.moco_v2 import MOCOV2Module
from pytorchvideo_trainer.module.simclr import SimCLRModule
from pytorchvideo_trainer.module.video_classification import VideoClassificationModule
class TestVideoClassificationModuleConf(unittest.TestCase):
def test_init_with_hydra(self) -> None:
with initialize_config_module(config_module="pytorchvideo_trainer.conf"):
test_conf = compose(
config_name="video_classification_train_app_conf",
overrides=["module/model=slow_r50"],
)
test_module = hydra.utils.instantiate(test_conf.module, _recursive_=False)
self.assertIsInstance(test_module, VideoClassificationModule)
self.assertIsNotNone(test_module.model)
class TestVideoSimCLRModuleConf(unittest.TestCase):
def test_init_with_hydra(self) -> None:
with initialize_config_module(config_module="pytorchvideo_trainer.conf"):
test_conf = compose(
config_name="simclr_train_app_conf",
)
test_module = hydra.utils.instantiate(test_conf.module, _recursive_=False)
self.assertIsInstance(test_module, SimCLRModule)
self.assertIsNotNone(test_module.model)
class TestVideoBYOLModuleConf(unittest.TestCase):
def test_init_with_hydra(self) -> None:
with initialize_config_module(config_module="pytorchvideo_trainer.conf"):
test_conf = compose(
config_name="byol_train_app_conf",
)
test_module = hydra.utils.instantiate(test_conf.module, _recursive_=False)
self.assertIsInstance(test_module, BYOLModule)
self.assertIsNotNone(test_module.model)
class TestVideoMOCOV2ModuleConf(unittest.TestCase):
def test_init_with_hydra(self) -> None:
with initialize_config_module(config_module="pytorchvideo_trainer.conf"):
test_conf = compose(
config_name="moco_v2_train_app_conf",
# overrides=["module/model=resnet"],
)
test_module = hydra.utils.instantiate(test_conf.module, _recursive_=False)
self.assertIsInstance(test_module, MOCOV2Module)
self.assertIsNotNone(test_module.model)
|