diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 705bba35..13f48107 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -94,10 +94,16 @@ jobs: - name: run tests bypassed run: PYTHONPATH=. python _unittests/ut_torch_models/test_tiny_llms_bypassed.py + - name: test image_classification + run: PYTHONPATH=. python _unittests/ut_tasks/test_tasks_image_classification.py + + - name: test zero_shot_image_classification + run: PYTHONPATH=. python _unittests/ut_tasks/test_tasks_zero_shot_image_classification.py + - name: run tests run: | pip install pytest - PYTHONPATH=. UNITTEST_GOING=1 pytest --durations=10 _unittests --ignore _unittests/ut_reference/test_backend_extended_reference_evaluator.py --ignore _unittests/ut_reference/test_backend_onnxruntime_evaluator.py --ignore _unittests/ut_torch_models/test_tiny_llms_bypassed.py + PYTHONPATH=. UNITTEST_GOING=1 pytest --durations=10 _unittests --ignore _unittests/ut_reference/test_backend_extended_reference_evaluator.py --ignore _unittests/ut_reference/test_backend_onnxruntime_evaluator.py --ignore _unittests/ut_torch_models/test_tiny_llms_bypassed.py --ignore _unittests/ut_tasks/test_tasks_zero_shot_image_classification.py --ignore _unittests/ut_tasks/test_tasks_image_classification.py - name: run backend tests python run: PYTHONPATH=. UNITTEST_GOING=1 pytest --durations=10 _unittests/ut_reference/test_backend_extended_reference_evaluator.py diff --git a/CHANGELOGS.rst b/CHANGELOGS.rst index d6812406..4eb46bf3 100644 --- a/CHANGELOGS.rst +++ b/CHANGELOGS.rst @@ -4,6 +4,7 @@ Change Logs 0.4.1 +++++ +* :pr:`72`: fix change_dynamic_dimension for custom classes * :pr:`70`: support models options in command lines 0.4.0 diff --git a/_doc/api/tasks/index.rst b/_doc/api/tasks/index.rst index 952dbee2..bafedabf 100644 --- a/_doc/api/tasks/index.rst +++ b/_doc/api/tasks/index.rst @@ -9,8 +9,11 @@ All submodules contains the three following functions: * ``random_input_kwargs(config) -> kwargs, get_inputs``: produces values ``get_inputs`` can take to generate dummy inputs suitable for a model defined by its configuration -* ``get_inputs(model, config, *args, **kwargs) -> dict(inputs=..., dynamic_shapes=...)``: - generates the dummy inputs and dynamic shapes for a specific model and configuration. +* ``get_inputs(model, config, *args, add_second_input=False, **kwargs) -> dict(inputs=..., dynamic_shapes=...)``: + generates the dummy inputs and dynamic shapes for a specific model and configuration, + if ``add_second_input`` is True, the function should return a different set of inputs, + with different values for the dynamic dimension. This is usually better to + rely on the function as the dynamic dimensions may be correlated. For a specific task, you would write: diff --git a/_unittests/ut_export/test_dynamic_shapes.py b/_unittests/ut_export/test_dynamic_shapes.py index be8fb56c..a4a13de4 100644 --- a/_unittests/ut_export/test_dynamic_shapes.py +++ b/_unittests/ut_export/test_dynamic_shapes.py @@ -1,5 +1,6 @@ import unittest import torch +import transformers from onnx_diagnostic.ext_test_case import ExtTestCase, requires_transformers from onnx_diagnostic.helpers import string_type from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache @@ -742,6 +743,18 @@ def test_couple_input_ds_change_dynamic_dimensions_fixed(self): self.assertEqual((1, 5, 8), new_input["A"].shape) self.assertEqual((1, 50), new_input["B"].shape) + def test_couple_input_ds_change_dynamic_dimensions_dynamic_cache(self): + inst = CoupleInputsDynamicShapes( + (), + {"A": make_dynamic_cache([(torch.ones((2, 2, 2, 2)), torch.ones((2, 2, 2, 2)))])}, + {"A": [[{0: "batch", 2: "last"}], [{0: "batch", 2: "last"}]]}, + ) + with bypass_export_some_errors(patch_transformers=True): + new_inputs = inst.change_dynamic_dimensions() + self.assertIsInstance(new_inputs["A"], transformers.cache_utils.DynamicCache) + self.assertEqual((3, 2, 3, 2), new_inputs["A"].key_cache[0].shape) + self.assertEqual((3, 2, 3, 2), new_inputs["A"].value_cache[0].shape) + @requires_transformers("4.51") def test_dynamic_cache_replace_by_string(self): n_layers = 2 diff --git a/_unittests/ut_helpers/test_helper.py b/_unittests/ut_helpers/test_helper.py index 046da51e..42a0d395 100644 --- a/_unittests/ut_helpers/test_helper.py +++ b/_unittests/ut_helpers/test_helper.py @@ -39,6 +39,8 @@ dtype_to_tensor_dtype, ) from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache, make_encoder_decoder_cache +from onnx_diagnostic.torch_models.hghub.hub_api import get_pretrained_config + TFLOAT = onnx.TensorProto.FLOAT @@ -484,6 +486,11 @@ def test_flatten_encoder_decoder_cache(self): s = string_type(inputs) self.assertIn("EncoderDecoderCache", s) + def test_string_typeçconfig(self): + conf = get_pretrained_config("microsoft/phi-2") + s = string_type(conf) + self.assertStartsWith("PhiConfig(**{", s) + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/_unittests/ut_tasks/test_tasks.py b/_unittests/ut_tasks/test_tasks.py index 1ff3bdf9..2b1368db 100644 --- a/_unittests/ut_tasks/test_tasks.py +++ b/_unittests/ut_tasks/test_tasks.py @@ -1,6 +1,6 @@ import unittest import torch -from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout, has_transformers, has_torch +from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout, has_transformers from onnx_diagnostic.torch_models.hghub.model_inputs import get_untrained_model_with_inputs from onnx_diagnostic.torch_export_patches import bypass_export_some_errors from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str @@ -10,11 +10,27 @@ class TestTasks(ExtTestCase): @hide_stdout() def test_text2text_generation(self): mid = "sshleifer/tiny-marian-en-de" - data = get_untrained_model_with_inputs(mid, verbose=1) + data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True) + self.assertEqual(data["task"], "text2text-generation") self.assertIn((data["size"], data["n_weights"]), [(473928, 118482)]) model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"] raise unittest.SkipTest(f"not working for {mid!r}") model(**inputs) + model(**data["inputs2"]) + with bypass_export_some_errors(patch_transformers=True, verbose=10): + torch.export.export( + model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False + ) + + @hide_stdout() + def test_text_generation(self): + mid = "arnir0/Tiny-LLM" + data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True) + self.assertEqual(data["task"], "text-generation") + self.assertIn((data["size"], data["n_weights"]), [(51955968, 12988992)]) + model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"] + model(**inputs) + model(**data["inputs2"]) with bypass_export_some_errors(patch_transformers=True, verbose=10): torch.export.export( model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False @@ -23,9 +39,11 @@ def test_text2text_generation(self): @hide_stdout() def test_automatic_speech_recognition(self): mid = "openai/whisper-tiny" - data = get_untrained_model_with_inputs(mid, verbose=1) + data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True) + self.assertEqual(data["task"], "automatic-speech-recognition") self.assertIn((data["size"], data["n_weights"]), [(132115968, 33028992)]) model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"] + model(**data["inputs2"]) Dim = torch.export.Dim self.maxDiff = None self.assertIn("{0:Dim(batch),1:DYN(seq_length)}", self.string_type(ds)) @@ -90,27 +108,15 @@ def test_automatic_speech_recognition(self): model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False ) - @hide_stdout() - def test_imagetext2text_generation(self): - mid = "HuggingFaceM4/tiny-random-idefics" - data = get_untrained_model_with_inputs(mid, verbose=1) - self.assertIn((data["size"], data["n_weights"]), [(12742888, 3185722)]) - model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"] - model(**inputs) - if not has_torch("2.10"): - raise unittest.SkipTest("sym_max does not work with dynamic dimension") - with bypass_export_some_errors(patch_transformers=True, verbose=10): - torch.export.export( - model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False - ) - @hide_stdout() def test_fill_mask(self): mid = "google-bert/bert-base-multilingual-cased" - data = get_untrained_model_with_inputs(mid, verbose=1) + data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True) + self.assertEqual(data["task"], "fill-mask") self.assertIn((data["size"], data["n_weights"]), [(428383212, 107095803)]) model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"] model(**inputs) + model(**data["inputs2"]) with bypass_export_some_errors(patch_transformers=True, verbose=10): torch.export.export( model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False @@ -119,10 +125,12 @@ def test_fill_mask(self): @hide_stdout() def test_feature_extraction(self): mid = "facebook/bart-base" - data = get_untrained_model_with_inputs(mid, verbose=1) + data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True) + self.assertEqual(data["task"], "feature-extraction") self.assertIn((data["size"], data["n_weights"]), [(557681664, 139420416)]) model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"] model(**inputs) + model(**data["inputs2"]) with bypass_export_some_errors(patch_transformers=True, verbose=10): torch.export.export( model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False @@ -131,10 +139,12 @@ def test_feature_extraction(self): @hide_stdout() def test_text_classification(self): mid = "Intel/bert-base-uncased-mrpc" - data = get_untrained_model_with_inputs(mid, verbose=1) + data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True) + self.assertEqual(data["task"], "text-classification") self.assertIn((data["size"], data["n_weights"]), [(154420232, 38605058)]) model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"] model(**inputs) + model(**data["inputs2"]) with bypass_export_some_errors(patch_transformers=True, verbose=10): torch.export.export( model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False @@ -143,10 +153,12 @@ def test_text_classification(self): @hide_stdout() def test_sentence_similary(self): mid = "sentence-transformers/all-MiniLM-L6-v1" - data = get_untrained_model_with_inputs(mid, verbose=1) + data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True) + self.assertEqual(data["task"], "sentence-similarity") self.assertIn((data["size"], data["n_weights"]), [(62461440, 15615360)]) model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"] model(**inputs) + model(**data["inputs2"]) with bypass_export_some_errors(patch_transformers=True, verbose=10): torch.export.export( model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False @@ -155,9 +167,11 @@ def test_sentence_similary(self): @hide_stdout() def test_falcon_mamba_dev(self): mid = "tiiuae/falcon-mamba-tiny-dev" - data = get_untrained_model_with_inputs(mid, verbose=1) + data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True) + self.assertEqual(data["task"], "text-generation") model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"] model(**inputs) + model(**data["inputs2"]) self.assertIn((data["size"], data["n_weights"]), [(138640384, 34660096)]) if not has_transformers("4.55"): raise unittest.SkipTest("The model has control flow.") diff --git a/_unittests/ut_tasks/test_tasks_image_classification.py b/_unittests/ut_tasks/test_tasks_image_classification.py new file mode 100644 index 00000000..bfb87332 --- /dev/null +++ b/_unittests/ut_tasks/test_tasks_image_classification.py @@ -0,0 +1,28 @@ +import unittest +import torch +from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout, has_transformers +from onnx_diagnostic.torch_models.hghub.model_inputs import get_untrained_model_with_inputs +from onnx_diagnostic.torch_export_patches import bypass_export_some_errors +from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str + + +class TestTasks(ExtTestCase): + @hide_stdout() + def test_image_classification(self): + mid = "hf-internal-testing/tiny-random-BeitForImageClassification" + data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True) + self.assertEqual(data["task"], "image-classification") + self.assertIn((data["size"], data["n_weights"]), [(56880, 14220)]) + model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"] + model(**inputs) + model(**data["inputs2"]) + if not has_transformers("4.51.999"): + raise unittest.SkipTest("Requires transformers>=4.52") + with bypass_export_some_errors(patch_transformers=True, verbose=10): + torch.export.export( + model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False + ) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/_unittests/ut_tasks/test_tasks_image_text_to_text.py b/_unittests/ut_tasks/test_tasks_image_text_to_text.py new file mode 100644 index 00000000..a99db3f4 --- /dev/null +++ b/_unittests/ut_tasks/test_tasks_image_text_to_text.py @@ -0,0 +1,30 @@ +import unittest +import torch +from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout, has_transformers, has_torch +from onnx_diagnostic.torch_models.hghub.model_inputs import get_untrained_model_with_inputs +from onnx_diagnostic.torch_export_patches import bypass_export_some_errors +from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str + + +class TestTasks(ExtTestCase): + @hide_stdout() + def test_image_text_to_text(self): + mid = "HuggingFaceM4/tiny-random-idefics" + data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True) + self.assertEqual(data["task"], "image-text-to-text") + self.assertIn((data["size"], data["n_weights"]), [(12742888, 3185722)]) + model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"] + model(**inputs) + model(**data["inputs2"]) + if not has_transformers("4.55"): + raise unittest.SkipTest("The model has control flow.") + if not has_torch("2.7.99"): + raise unittest.SkipTest("sym_max does not work with dynamic dimension") + with bypass_export_some_errors(patch_transformers=True, verbose=10): + torch.export.export( + model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False + ) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/_unittests/ut_tasks/test_tasks_zero_shot_image_classification.py b/_unittests/ut_tasks/test_tasks_zero_shot_image_classification.py new file mode 100644 index 00000000..b441646a --- /dev/null +++ b/_unittests/ut_tasks/test_tasks_zero_shot_image_classification.py @@ -0,0 +1,27 @@ +import unittest +import torch +from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout, requires_torch +from onnx_diagnostic.torch_models.hghub.model_inputs import get_untrained_model_with_inputs +from onnx_diagnostic.torch_export_patches import bypass_export_some_errors +from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str + + +class TestTasks(ExtTestCase): + @requires_torch("2.7.99") + @hide_stdout() + def test_zero_shot_image_classification(self): + mid = "openai/clip-vit-base-patch16" + data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True) + self.assertEqual(data["task"], "zero-shot-image-classification") + self.assertIn((data["size"], data["n_weights"]), [(188872708, 47218177)]) + model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"] + model(**inputs) + model(**data["inputs2"]) + with bypass_export_some_errors(patch_transformers=True, verbose=10): + torch.export.export( + model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False + ) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/onnx_diagnostic/export/dynamic_shapes.py b/onnx_diagnostic/export/dynamic_shapes.py index 28fa5f00..a89a9e68 100644 --- a/onnx_diagnostic/export/dynamic_shapes.py +++ b/onnx_diagnostic/export/dynamic_shapes.py @@ -363,16 +363,20 @@ def _generic_walker_step( ) if flatten_unflatten: flatunflat = flatten_unflatten_for_dynamic_shapes(inputs) - return cls._generic_walker_step( + res = cls._generic_walker_step( processor, flatunflat, ds, flatten_unflatten=flatten_unflatten ) - flat, _spec = torch.utils._pytree.tree_flatten(inputs) + # Should we restore the original class? + return res + flat, spec = torch.utils._pytree.tree_flatten(inputs) if all(isinstance(t, torch.Tensor) for t in flat): # We need to flatten dynamic shapes as well ds = flatten_dynamic_shapes(ds) - return cls._generic_walker_step( + res = cls._generic_walker_step( processor, flat, ds, flatten_unflatten=flatten_unflatten ) + # Then we restore the original class. + return torch.utils._pytree.tree_unflatten(res, spec) class ChangeDimensionProcessor: def __init__(self, desired_values): diff --git a/onnx_diagnostic/helpers/helper.py b/onnx_diagnostic/helpers/helper.py index ae6379f1..b61d0a48 100644 --- a/onnx_diagnostic/helpers/helper.py +++ b/onnx_diagnostic/helpers/helper.py @@ -666,6 +666,15 @@ def string_type( print(f"[string_type] CACHE4:{type(obj)}") return f"{obj.__class__.__name__}(...)" + if obj.__class__.__name__.endswith("Config"): + import transformers.configuration_utils as tcu + + if isinstance(obj, tcu.PretrainedConfig): + if verbose: + print(f"[string_type] CONFIG:{type(obj)}") + s = str(obj.to_diff_dict()).replace("\n", "").replace(" ", "") + return f"{obj.__class__.__name__}(**{s})" + if verbose: print(f"[string_type] END:{type(obj)}") raise AssertionError(f"Unsupported type {type(obj).__name__!r} - {type(obj)}") diff --git a/onnx_diagnostic/tasks/automatic_speech_recognition.py b/onnx_diagnostic/tasks/automatic_speech_recognition.py index d5ec5475..346c3fa2 100644 --- a/onnx_diagnostic/tasks/automatic_speech_recognition.py +++ b/onnx_diagnostic/tasks/automatic_speech_recognition.py @@ -33,6 +33,7 @@ def get_inputs( head_dim: int, batch_size: int = 2, sequence_length: int = 30, + add_second_input: bool = False, **kwargs, # unused ): """ @@ -126,7 +127,24 @@ def get_inputs( # encoder_last_hidden_state=torch.randn(batch_size, sequence_length2, encoder_dim), # encoder_outputs=torch.randn(batch_size, sequence_length2, encoder_dim), ) - return dict(inputs=inputs, dynamic_shapes=shapes) + res = dict(inputs=inputs, dynamic_shapes=shapes) + if add_second_input: + res["inputs2"] = get_inputs( + model=model, + config=config, + dummy_max_token_id=dummy_max_token_id, + max_source_positions=max_source_positions, + d_model=d_model, + num_hidden_layers=num_hidden_layers, + encoder_attention_heads=encoder_attention_heads, + encoder_layers=encoder_layers, + decoder_layers=decoder_layers, + head_dim=head_dim, + batch_size=batch_size + 1, + sequence_length=sequence_length + 1, + **kwargs, + )["inputs"] + return res def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]: diff --git a/onnx_diagnostic/tasks/feature_extraction.py b/onnx_diagnostic/tasks/feature_extraction.py index 510a9f1f..9ef52058 100644 --- a/onnx_diagnostic/tasks/feature_extraction.py +++ b/onnx_diagnostic/tasks/feature_extraction.py @@ -22,6 +22,7 @@ def get_inputs( batch_size: int, sequence_length: int, dummy_max_token_id: int, + add_second_input: bool = False, **kwargs, # unused ): """ @@ -46,7 +47,17 @@ def get_inputs( ), attention_mask=torch.ones((batch_size, sequence_length)).to(torch.int64), ) - return dict(inputs=inputs, dynamic_shapes=shapes) + res = dict(inputs=inputs, dynamic_shapes=shapes) + if add_second_input: + res["inputs2"] = get_inputs( + model=model, + config=config, + batch_size=batch_size + 1, + sequence_length=sequence_length + 1, + dummy_max_token_id=dummy_max_token_id, + **kwargs, + )["inputs"] + return res def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]: diff --git a/onnx_diagnostic/tasks/fill_mask.py b/onnx_diagnostic/tasks/fill_mask.py index 5f0f1af8..14020e1e 100644 --- a/onnx_diagnostic/tasks/fill_mask.py +++ b/onnx_diagnostic/tasks/fill_mask.py @@ -22,6 +22,7 @@ def get_inputs( batch_size: int, sequence_length: int, dummy_max_token_id: int, + add_second_input: bool = False, **kwargs, # unused ): """ @@ -48,7 +49,17 @@ def get_inputs( token_type_ids=torch.zeros((batch_size, sequence_length)).to(torch.int64), attention_mask=torch.ones((batch_size, sequence_length)).to(torch.int64), ) - return dict(inputs=inputs, dynamic_shapes=shapes) + res = dict(inputs=inputs, dynamic_shapes=shapes) + if add_second_input: + res["inputs2"] = get_inputs( + model=model, + config=config, + batch_size=batch_size + 1, + sequence_length=sequence_length + 1, + dummy_max_token_id=dummy_max_token_id, + **kwargs, + )["inputs"] + return res def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]: diff --git a/onnx_diagnostic/tasks/image_classification.py b/onnx_diagnostic/tasks/image_classification.py index d9876936..2d0696f2 100644 --- a/onnx_diagnostic/tasks/image_classification.py +++ b/onnx_diagnostic/tasks/image_classification.py @@ -27,6 +27,7 @@ def get_inputs( input_channels: int, batch_size: int = 2, dynamic_rope: bool = False, + add_second_input: bool = False, **kwargs, # unused ): """ @@ -59,7 +60,19 @@ def get_inputs( -1, 1 ), ) - return dict(inputs=inputs, dynamic_shapes=shapes) + res = dict(inputs=inputs, dynamic_shapes=shapes) + if add_second_input: + res["inputs2"] = get_inputs( + model=model, + config=config, + input_width=input_width + 1, + input_height=input_height + 1, + input_channels=input_channels, + batch_size=batch_size + 1, + dynamic_rope=dynamic_rope, + **kwargs, + )["inputs"] + return res def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]: diff --git a/onnx_diagnostic/tasks/image_text_to_text.py b/onnx_diagnostic/tasks/image_text_to_text.py index 2e4e5eae..50014621 100644 --- a/onnx_diagnostic/tasks/image_text_to_text.py +++ b/onnx_diagnostic/tasks/image_text_to_text.py @@ -32,6 +32,7 @@ def get_inputs( sequence_length2: int = 3, n_images: int = 2, dynamic_rope: bool = False, + add_second_input: bool = False, **kwargs, # unused ): """ @@ -99,7 +100,26 @@ def get_inputs( torch.int64 ), ) - return dict(inputs=inputs, dynamic_shapes=shapes) + res = dict(inputs=inputs, dynamic_shapes=shapes) + if add_second_input: + res["inputs2"] = get_inputs( + model=model, + config=config, + dummy_max_token_id=dummy_max_token_id, + num_key_value_heads=num_key_value_heads, + num_hidden_layers=num_hidden_layers, + head_dim=head_dim, + width=width, + height=height, + num_channels=num_channels, + batch_size=batch_size + 1, + sequence_length=sequence_length + 1, + sequence_length2=sequence_length2 + 1, + n_images=n_images + 1, + dynamic_rope=dynamic_rope, + **kwargs, + )["inputs"] + return res def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]: diff --git a/onnx_diagnostic/tasks/mixture_of_expert.py b/onnx_diagnostic/tasks/mixture_of_expert.py index d7d302ac..b7e5af37 100644 --- a/onnx_diagnostic/tasks/mixture_of_expert.py +++ b/onnx_diagnostic/tasks/mixture_of_expert.py @@ -41,6 +41,7 @@ def get_inputs( sequence_length2: int = 3, n_images: int = 2, dynamic_rope: bool = False, + add_second_input: bool = False, **kwargs, # unused ): """ @@ -60,6 +61,7 @@ def get_inputs( :param dynamic_rope: use dynamic rope (see :class:`transformers.LlamaConfig`) :return: dictionary """ + assert not add_second_input, "add_second_input=True not yet implemented" raise NotImplementedError(f"get_inputs not yet implemented for task {__TASK__!r}.") diff --git a/onnx_diagnostic/tasks/sentence_similarity.py b/onnx_diagnostic/tasks/sentence_similarity.py index df9f00e3..808ae039 100644 --- a/onnx_diagnostic/tasks/sentence_similarity.py +++ b/onnx_diagnostic/tasks/sentence_similarity.py @@ -22,6 +22,7 @@ def get_inputs( batch_size: int, sequence_length: int, dummy_max_token_id: int, + add_second_input: bool = False, **kwargs, # unused ): """ @@ -48,7 +49,17 @@ def get_inputs( token_type_ids=torch.zeros((batch_size, sequence_length)).to(torch.int64), attention_mask=torch.ones((batch_size, sequence_length)).to(torch.int64), ) - return dict(inputs=inputs, dynamic_shapes=shapes) + res = dict(inputs=inputs, dynamic_shapes=shapes) + if add_second_input: + res["inputs2"] = get_inputs( + model=model, + config=config, + batch_size=batch_size + 1, + sequence_length=sequence_length + 1, + dummy_max_token_id=dummy_max_token_id, + **kwargs, + )["inputs"] + return res def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]: diff --git a/onnx_diagnostic/tasks/text2text_generation.py b/onnx_diagnostic/tasks/text2text_generation.py index 02edf654..dde8841c 100644 --- a/onnx_diagnostic/tasks/text2text_generation.py +++ b/onnx_diagnostic/tasks/text2text_generation.py @@ -28,6 +28,7 @@ def get_inputs( batch_size: int = 2, sequence_length: int = 30, sequence_length2: int = 3, + add_second_input: bool = False, **kwargs, # unused ): """ @@ -125,7 +126,22 @@ def get_inputs( # encoder_last_hidden_state=torch.randn(batch_size, sequence_length2, encoder_dim), # encoder_outputs=torch.randn(batch_size, sequence_length2, encoder_dim), ) - return dict(inputs=inputs, dynamic_shapes=shapes) + res = dict(inputs=inputs, dynamic_shapes=shapes) + if add_second_input: + res["inputs2"] = get_inputs( + model=model, + config=config, + dummy_max_token_id=dummy_max_token_id, + num_key_value_heads=num_key_value_heads, + num_hidden_layers=num_hidden_layers, + head_dim=head_dim, + encoder_dim=encoder_dim, + batch_size=batch_size + 1, + sequence_length=sequence_length + 1, + sequence_length2=sequence_length2 + 1, + **kwargs, + )["inputs"] + return res def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]: diff --git a/onnx_diagnostic/tasks/text_classification.py b/onnx_diagnostic/tasks/text_classification.py index 15a88d2a..aaaa8838 100644 --- a/onnx_diagnostic/tasks/text_classification.py +++ b/onnx_diagnostic/tasks/text_classification.py @@ -22,6 +22,7 @@ def get_inputs( batch_size: int, sequence_length: int, dummy_max_token_id: int, + add_second_input: bool = False, **kwargs, # unused ): """ @@ -48,7 +49,17 @@ def get_inputs( token_type_ids=torch.zeros((batch_size, sequence_length)).to(torch.int64), attention_mask=torch.ones((batch_size, sequence_length)).to(torch.int64), ) - return dict(inputs=inputs, dynamic_shapes=shapes) + res = dict(inputs=inputs, dynamic_shapes=shapes) + if add_second_input: + res["inputs2"] = get_inputs( + model=model, + config=config, + batch_size=batch_size + 1, + sequence_length=sequence_length + 1, + dummy_max_token_id=dummy_max_token_id, + **kwargs, + )["inputs"] + return res def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]: diff --git a/onnx_diagnostic/tasks/text_generation.py b/onnx_diagnostic/tasks/text_generation.py index dc6cf914..897c7a38 100644 --- a/onnx_diagnostic/tasks/text_generation.py +++ b/onnx_diagnostic/tasks/text_generation.py @@ -71,6 +71,7 @@ def get_inputs( num_key_value_heads: Optional[int] = None, head_dim: Optional[int] = None, cls_cache: Optional[Union[type, str]] = None, + add_second_input: bool = False, **kwargs, # unused ): """ @@ -144,55 +145,75 @@ def get_inputs( ] ), ) - return dict(inputs=inputs, dynamic_shapes=shapes) - - if head_dim is None: - assert config, "head_dim is None, the value cannot be set without a configuration" - head_dim = config.hidden_size // config.num_attention_heads + res = dict(inputs=inputs, dynamic_shapes=shapes) + else: + if head_dim is None: + assert config, "head_dim is None, the value cannot be set without a configuration" + head_dim = config.hidden_size // config.num_attention_heads - shapes = { - "input_ids": {0: batch, 1: seq_length}, - "attention_mask": { - 0: batch, - 1: "cache+seq", # cache_length + seq_length - }, - "position_ids": { - 0: batch, - 1: "cache+seq", # cache_length + seq_length - }, - "past_key_values": [ - [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)], - [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)], - ], - } + shapes = { + "input_ids": {0: batch, 1: seq_length}, + "attention_mask": { + 0: batch, + 1: "cache+seq", # cache_length + seq_length + }, + "position_ids": { + 0: batch, + 1: "cache+seq", # cache_length + seq_length + }, + "past_key_values": [ + [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)], + [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)], + ], + } - make_cache = ( - make_sliding_window_cache - if cls_cache in ("SlidingWindowCache", transformers.cache_utils.SlidingWindowCache) - else make_dynamic_cache - ) + make_cache = ( + make_sliding_window_cache + if cls_cache in ("SlidingWindowCache", transformers.cache_utils.SlidingWindowCache) + else make_dynamic_cache + ) - inputs = dict( - input_ids=torch.randint(0, dummy_max_token_id, (batch_size, sequence_length2)).to( - torch.int64 - ), - attention_mask=torch.ones((batch_size, sequence_length + sequence_length2)).to( - torch.int64 - ), - position_ids=torch.arange(sequence_length, sequence_length + sequence_length2) - .to(torch.int64) - .expand((batch_size, -1)), - past_key_values=make_cache( - [ - ( - torch.randn(batch_size, num_key_value_heads, sequence_length, head_dim), - torch.randn(batch_size, num_key_value_heads, sequence_length, head_dim), - ) - for i in range(num_hidden_layers) - ] - ), - ) - return dict(inputs=inputs, dynamic_shapes=shapes) + inputs = dict( + input_ids=torch.randint(0, dummy_max_token_id, (batch_size, sequence_length2)).to( + torch.int64 + ), + attention_mask=torch.ones((batch_size, sequence_length + sequence_length2)).to( + torch.int64 + ), + position_ids=torch.arange(sequence_length, sequence_length + sequence_length2) + .to(torch.int64) + .expand((batch_size, -1)), + past_key_values=make_cache( + [ + ( + torch.randn( + batch_size, num_key_value_heads, sequence_length, head_dim + ), + torch.randn( + batch_size, num_key_value_heads, sequence_length, head_dim + ), + ) + for i in range(num_hidden_layers) + ] + ), + ) + res = dict(inputs=inputs, dynamic_shapes=shapes) + if add_second_input: + res["inputs2"] = get_inputs( + model=model, + config=config, + dummy_max_token_id=dummy_max_token_id, + num_hidden_layers=num_hidden_layers, + batch_size=batch_size + 1, + sequence_length=sequence_length + 1, + sequence_length2=sequence_length2 + 1, + dynamic_rope=dynamic_rope, + num_key_value_heads=num_key_value_heads, + head_dim=head_dim, + cls_cache=cls_cache, + **kwargs, + )["inputs"] + return res def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]: diff --git a/onnx_diagnostic/tasks/zero_shot_image_classification.py b/onnx_diagnostic/tasks/zero_shot_image_classification.py index 42da5628..a341a191 100644 --- a/onnx_diagnostic/tasks/zero_shot_image_classification.py +++ b/onnx_diagnostic/tasks/zero_shot_image_classification.py @@ -34,6 +34,7 @@ def get_inputs( input_height: int = 224, input_channels: int = 3, batch_size_image=3, + add_second_input: bool = False, **kwargs, # unused ): """ @@ -81,7 +82,21 @@ def get_inputs( batch_size_image, input_channels, input_width, input_height ).clamp(-1, 1), ) - return dict(inputs=inputs, dynamic_shapes=shapes) + res = dict(inputs=inputs, dynamic_shapes=shapes) + if add_second_input: + res["inputs2"] = get_inputs( + model=model, + config=config, + dummy_max_token_id=dummy_max_token_id, + batch_size=batch_size + 1, + sequence_length=sequence_length + 1, + input_width=input_width, + input_height=input_height, + input_channels=input_channels, + batch_size_image=batch_size_image + 1, + **kwargs, + )["inputs"] + return res def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]: diff --git a/onnx_diagnostic/torch_export_patches/patches/patch_torch.py b/onnx_diagnostic/torch_export_patches/patches/patch_torch.py index 0df84c00..d53e3cc9 100644 --- a/onnx_diagnostic/torch_export_patches/patches/patch_torch.py +++ b/onnx_diagnostic/torch_export_patches/patches/patch_torch.py @@ -44,7 +44,7 @@ def _catch_produce_guards_and_solve_constraints( raise if verbose: print( - f"[_catch_produce_guards_and_solve_constraints] ERROR" + f"[_catch_produce_guards_and_solve_constraints] ERROR: " f"produce_guards_and_solve_constraints failed, " f"use SKIP_SOLVE_CONSTRAINTS=0 to avoid skipping\n" f"fake_mode={fake_mode}\n" @@ -54,6 +54,7 @@ def _catch_produce_guards_and_solve_constraints( f"_is_torch_jit_trace={_is_torch_jit_trace}\n" f"exc={e}\ngm={gm}" ) + torch._dynamo.reset() def patch__check_input_constraints_for_graph( @@ -70,13 +71,14 @@ def patch__check_input_constraints_for_graph( raise if verbose: print( - f"[_check_input_constraints_for_graph] ERROR" + f"[_check_input_constraints_for_graph] ERROR: " f"_check_input_constraints_for_graph failed, " f"use SKIP_SOLVE_CONSTRAINTS=0 to avoid skipping\n" f"input_placeholders={input_placeholders}\n" f"range_constraints={range_constraints}\n" f"exc={e}" ) + torch._dynamo.reset() def patched_infer_size(a, b): diff --git a/onnx_diagnostic/torch_export_patches/patches/patch_transformers.py b/onnx_diagnostic/torch_export_patches/patches/patch_transformers.py index e45b211a..a27005ce 100644 --- a/onnx_diagnostic/torch_export_patches/patches/patch_transformers.py +++ b/onnx_diagnostic/torch_export_patches/patches/patch_transformers.py @@ -5,6 +5,7 @@ import transformers from transformers.modeling_attn_mask_utils import AttentionMaskConverter from transformers.cache_utils import StaticCache, Cache, DynamicCache +from ...ext_test_case import has_transformers from ...helpers.torch_test_helper import is_torchdynamo_exporting @@ -50,7 +51,8 @@ class patched_AttentionMaskConverter: ``transformers.modeling_attn_mask_utils.AttentionMaskConverter._make_causal_mask``. """ - _PATCHES_ = ["_make_causal_mask"] + # This method was fixed in 4.51 at least. + _PATCHES_ = ["_make_causal_mask"] if not has_transformers("4.48.3") else [] _PATCHED_CLASS_ = AttentionMaskConverter @staticmethod @@ -69,6 +71,9 @@ def _make_causal_mask( This static method may be called with ``AttentionMaskConverter._make_causal_mask`` or ``self._make_causal_mask``. That changes this argument is receives. That should not matter but... + The patch should be implemented in another way. static methods do not play well + with a simple replacement. + Fortunately, this patch does not seem to be needed anymore with transformers>=4.48.3. """ if args: index = 0 if isinstance(args[0], (tuple, torch.Size)) else 1 diff --git a/onnx_diagnostic/torch_models/hghub/hub_data_cached_configs.py b/onnx_diagnostic/torch_models/hghub/hub_data_cached_configs.py index d69b23cd..4e48a503 100644 --- a/onnx_diagnostic/torch_models/hghub/hub_data_cached_configs.py +++ b/onnx_diagnostic/torch_models/hghub/hub_data_cached_configs.py @@ -3569,3 +3569,74 @@ def _ccached_tiiuae_falcon_mamba_tiny_dev(): "vocab_size": 65024, } ) + + +def _ccached_facebook_bart_base(): + "facebook/bart-base" + return transformers.BartConfig( + **{ + "_name_or_path": "bart-base", + "activation_dropout": 0.1, + "activation_function": "gelu", + "add_bias_logits": false, + "add_final_layer_norm": false, + "architectures": ["BartModel"], + "attention_dropout": 0.1, + "bos_token_id": 0, + "classif_dropout": 0.1, + "classifier_dropout": 0.0, + "d_model": 768, + "decoder_attention_heads": 12, + "decoder_ffn_dim": 3072, + "decoder_layerdrop": 0.0, + "decoder_layers": 6, + "decoder_start_token_id": 2, + "dropout": 0.1, + "early_stopping": true, + "encoder_attention_heads": 12, + "encoder_ffn_dim": 3072, + "encoder_layerdrop": 0.0, + "encoder_layers": 6, + "eos_token_id": 2, + "forced_eos_token_id": 2, + "forced_bos_token_id": 0, + "gradient_checkpointing": false, + "id2label": {"0": "LABEL_0", "1": "LABEL_1", "2": "LABEL_2"}, + "init_std": 0.02, + "is_encoder_decoder": true, + "label2id": {"LABEL_0": 0, "LABEL_1": 1, "LABEL_2": 2}, + "max_position_embeddings": 1024, + "model_type": "bart", + "no_repeat_ngram_size": 3, + "normalize_before": false, + "normalize_embedding": true, + "num_beams": 4, + "num_hidden_layers": 6, + "pad_token_id": 1, + "scale_embedding": false, + "task_specific_params": { + "summarization": { + "length_penalty": 1.0, + "max_length": 128, + "min_length": 12, + "num_beams": 4, + }, + "summarization_cnn": { + "length_penalty": 2.0, + "max_length": 142, + "min_length": 56, + "num_beams": 4, + }, + "summarization_xsum": { + "length_penalty": 1.0, + "max_length": 62, + "min_length": 11, + "num_beams": 6, + }, + }, + "torch_dtype": "float32", + "transformers_version": "4.12.0.dev0", + "use_cache": true, + "vocab_size": 50265, + } + ) diff --git a/onnx_diagnostic/torch_models/hghub/model_inputs.py b/onnx_diagnostic/torch_models/hghub/model_inputs.py index 6a1c6bf8..f4ab2650 100644 --- a/onnx_diagnostic/torch_models/hghub/model_inputs.py +++ b/onnx_diagnostic/torch_models/hghub/model_inputs.py @@ -17,6 +17,7 @@ def get_untrained_model_with_inputs( dynamic_rope: Optional[bool] = None, same_as_pretrained: bool = False, use_preinstalled: bool = True, + add_second_input: bool = False, ) -> Dict[str, Any]: """ Gets a non initialized model similar to the original model @@ -34,6 +35,8 @@ def get_untrained_model_with_inputs( :param same_as_pretrained: if True, do not change the default values to get a smaller model :param use_preinstalled: use preinstalled configurations + :param add_second_input: provides a second inputs to check a model + supports different shapes :return: dictionary with a model, inputs, dynamic shapes, and the configuration Example: @@ -107,7 +110,7 @@ def get_untrained_model_with_inputs( # This line is important. Some models may produce different # outputs even with the same inputs in training mode. model.eval() - res = fct(model, config, **kwargs) + res = fct(model, config, add_second_input=add_second_input, **kwargs) res["input_kwargs"] = kwargs res["model_kwargs"] = mkwargs