github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-05 11:02:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Analyzer] fix analyzer tests (#3197)

Browse Source
This commit is contained in:
YuliangLiu0306
2023-03-22 13:38:11 +08:00
committed by GitHub
parent f57d34958b
commit 019a847432
7 changed files with 60 additions and 100 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
tests/test_analyzer/test_fx/test_shape_prop.py
View File

@@ -1,16 +1,17 @@
import pytest
import timm.models as tmm
import torch
import torchvision.models as tm
from .zoo import tm_models, tmm_models
from packaging import version
from colossalai.testing.utils import parameterize
from tests.test_analyzer.test_fx.zoo import tm_models, tmm_models
try:
from colossalai._analyzer._subclasses import MetaTensorMode
from colossalai._analyzer.fx import symbolic_trace
from colossalai._analyzer.fx.passes.shape_prop import shape_prop_pass
from colossalai._analyzer.fx.symbolic_profile import register_shape_impl
@register_shape_impl(torch.nn.functional.linear)
def linear_impl(*args, **kwargs):
assert True
@@ -23,15 +24,15 @@ def _check_gm_validity(gm: torch.fx.GraphModule):
for node in gm.graph.nodes:
assert node.meta['info'].outputs, f'In {gm.__class__.__name__}, {node} has no output shape.'
if node.op in [
# 'call_module', # can apply to params
# 'call_function', # can apply to params
# 'call_method', # can apply to params
'call_module', # can apply to params
'call_function', # can apply to params
'call_method', # can apply to params
]:
assert node.meta['info'].inputs, f'In {gm.__class__.__name__}, {node} has no input shape.'
assert hasattr(node.meta['info'], 'inputs'), f'In {gm.__class__.__name__}, {node} has no input shape.'
@pytest.mark.skipif(torch.__version__ < '1.12.0', reason='torch version < 12')
@pytest.mark.parametrize('m', tm_models)
@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason='torch version < 12')
@parameterize('m', tm_models)
def test_torchvision_shape_prop(m):
with MetaTensorMode():
model = m()
@@ -44,8 +45,8 @@ def test_torchvision_shape_prop(m):
_check_gm_validity(gm)
@pytest.mark.skipif(torch.__version__ < '1.12.0', reason='torch version < 12')
@pytest.mark.parametrize('m', tmm_models)
@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.12.0'), reason='torch version < 12')
@parameterize('m', tmm_models)
def test_timm_shape_prop(m):
with MetaTensorMode():
model = m()
@@ -53,11 +54,12 @@ def test_timm_shape_prop(m):
meta_args = {
"x": data,
}
gm = symbolic_trace(model, meta_args=meta_args)
shape_prop_pass(gm, data)
_check_gm_validity(gm)
if __name__ == "__main__":
test_torchvision_shape_prop(tm.resnet18)
test_timm_shape_prop(tmm.vgg11)
test_torchvision_shape_prop()
test_timm_shape_prop()