2023-03-17 05:49:04 +00:00
|
|
|
import pytest
|
2023-05-23 03:57:15 +00:00
|
|
|
from lazy_init_utils import SUPPORT_LAZY, check_lazy_init
|
2023-03-17 05:49:04 +00:00
|
|
|
|
|
|
|
from tests.kit.model_zoo import model_zoo
|
|
|
|
|
|
|
|
|
2023-05-15 09:20:56 +00:00
|
|
|
@pytest.mark.skipif(not SUPPORT_LAZY, reason='requires torch >= 1.12.0')
|
2023-03-17 05:49:04 +00:00
|
|
|
@pytest.mark.parametrize('subset', ['torchvision', 'diffusers', 'timm', 'transformers', 'torchaudio', 'deepfm', 'dlrm'])
|
2023-07-19 08:43:01 +00:00
|
|
|
@pytest.mark.parametrize('default_device', ['cpu', 'cuda'])
|
|
|
|
def test_torchvision_models_lazy_init(subset, default_device):
|
2023-03-17 05:49:04 +00:00
|
|
|
sub_model_zoo = model_zoo.get_sub_registry(subset)
|
|
|
|
for name, entry in sub_model_zoo.items():
|
|
|
|
# TODO(ver217): lazy init does not support weight norm, skip these models
|
2023-06-26 07:50:07 +00:00
|
|
|
if name in ('torchaudio_wav2vec2_base', 'torchaudio_hubert_base') or name.startswith('transformers_llama'):
|
2023-03-17 05:49:04 +00:00
|
|
|
continue
|
2023-07-19 08:43:01 +00:00
|
|
|
check_lazy_init(entry, verbose=True, default_device=default_device)
|
2023-03-17 05:49:04 +00:00
|
|
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
test_torchvision_models_lazy_init('torchvision')
|