mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-08 20:40:34 +00:00
[lazy] refactor lazy init (#3891)
* [lazy] remove old lazy init * [lazy] refactor lazy init folder structure * [lazy] fix lazy tensor deepcopy * [test] update lazy init test
This commit is contained in:
19
tests/test_lazy/test_models.py
Normal file
19
tests/test_lazy/test_models.py
Normal file
@@ -0,0 +1,19 @@
|
||||
import pytest
|
||||
from lazy_init_utils import SUPPORT_LAZY, check_lazy_init
|
||||
|
||||
from tests.kit.model_zoo import model_zoo
|
||||
|
||||
|
||||
@pytest.mark.skipif(not SUPPORT_LAZY, reason='requires torch >= 1.12.0')
|
||||
@pytest.mark.parametrize('subset', ['torchvision', 'diffusers', 'timm', 'transformers', 'torchaudio', 'deepfm', 'dlrm'])
|
||||
def test_torchvision_models_lazy_init(subset):
|
||||
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
|
||||
if name in ('torchaudio_wav2vec2_base', 'torchaudio_hubert_base'):
|
||||
continue
|
||||
check_lazy_init(entry, verbose=True)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_torchvision_models_lazy_init('torchvision')
|
Reference in New Issue
Block a user