Graph prepare refactor by reuvenperetz · Pull Request #1490 · SonySemiconductorSolutions/mct-model-optimization

reuvenperetz · 2025-07-02T13:25:47Z

Pull Request Description:

Extract the parts of graph building from the framework model to an outer package called "graph_builder".
BaseGraphBuilder defines API methods for converting the model to a graph and transforming it using the fundamental substitutions.

Checklist before requesting a review:

I set the appropriate labels on the pull request.
I have added/updated the release note draft (if necessary).
I have updated the documentation to reflect my changes (if necessary).
All function and files are well documented.
All function and classes have type hints.
There is a licenses in all file.
The function and variable names are informative.
I have checked for code duplications.
I have added new unittest (if necessary).

…tq facade

…raph

…ide of the core runner

…zed_graph checks

ofirgo · 2025-07-03T07:07:54Z

@@ -0,0 +1,14 @@
+# Copyright 2025 Sony Semiconductor Israel, Inc. All rights reserved.


I think that graph_preparation is a better name for the package

ofirgo · 2025-07-03T07:12:16Z

+        if tensorboard_writer is not None:
+            tensorboard_writer.add_graph(graph, 'initial_graph')
+
+        if fqc:


I don't think that fqc should be an argument to this module, since it is quantization-related.
If we must still store it in the graph, then it supposed to be the "quantization_preparation" responsibility

I started without it, but the relu_to_pot substitution requires it, which is why I added it eventually.
I'll do a different PR to remove this requirement first.

ofirgo · 2025-07-03T07:15:52Z

+                    model: Any,
+                    representative_dataset: Callable = None,
+                    fqc: FrameworkQuantizationCapabilities = None,
+                    tensorboard_writer: TensorboardWriter = None,


Several things regarding the arguments:

Move tensor board to be the last argument.

Consider maybe gathering the 3 "structure-altering" flags in a dedicated config?

Do we need the default values to not "break" API/tests? if not, then remove the default.

Change the typehint of arguments that accept None to Optional[]

About 2: That's how I started, but I eventually thought it would cause unnecessary maintenance. I can write it back.

So depends on what the unnecessary maintanance is.
Normally, I wouldn't suggest to just have a config as an argument that is immidiatly flattend in the class init, but in this case they are all substitution bollean flags so it feels right to tie them together in a "which substitutions to run" config variable

ofirgo · 2025-07-03T07:17:59Z

+import torch
+
+
+def convert_pytorch_model_to_graph(model: torch.nn.Module,


Why split it to a separate file and not include the implementation as part of PytorchGraphBuilder's convert_model_to_graph method? (same for keras)

ofirgo · 2025-07-03T07:19:07Z

+from model_compression_toolkit.core.common.substitutions.apply_substitutions import substitute
+
+
+def transform_pytorch_graph(graph: Graph,


same question, why not as part of the pytorch builder class?

ofirgo · 2025-07-03T07:23:35Z

                                                                target_resource_utilization=target_resource_utilization,
-                                                                tb_w=tb_w)
+                                                                tb_w=tb_w,
+                                                                fw_graph_builder=KerasGraphBuilder())


Please add a todo either here (in each facade) or in the core_runner that this should eventually be initialized and run outside the runner.

I was trying to avoid duplicating code. Do you see any benefit in moving it outside rather than passing the framework class builder?

Eventually, the runner is going to die. The point of separating each logical block from the core into a new module is exactly this - to not have 1 backbone that you have to run in order to get through anything in the MCT.
So, the endgame is to have each facade call the 4-5 modules that it needs to - construct graph, put quantization info, run some optimization, run any enhancement that is part of the facade's algorithmic framework.

To summarize my long answer - yes, I see benefites in moving it outside of the runner

ofirgo · 2025-07-03T07:26:32Z

+from model_compression_toolkit.graph_builder.keras.convert_keras_model_to_graph import convert_keras_model_to_graph


 class TestGraphReading(unittest.TestCase):


As part of the PR where you impement tests for the new module (GraphBuilder), also migrate/remove this test and any other test that is meant to verify the graph building and transformation

ofirgo · 2025-07-03T09:40:32Z

+                    representative_dataset: Callable = None,
+                    fqc: FrameworkQuantizationCapabilities = None,
+                    tensorboard_writer: TensorboardWriter = None,
+                    linear_collapsing: bool = True,


what about the substitutions in get_substitutions_post_statistics_collection? aren't they graph preparation related? if so, need to think if they need to be part of this module or somewhere else

They depend on the quantization params (this is why they are used today only after loading the fqc, and computing quantization params). So I do not think they should be part of it...

Let's talk about this offline, I want us to take another look at this an figure out where they should land.

From a quick check this is what I think can work:

softmax shift - not quantization related, just need to happen after statistic collection, so maybe put it in this module but not part of the "graph_builder" flow (as a separate api that can be called)

concat threshold - this needs to happen after parameters selection, so it is like any other "quantization enhancement" substitution that we have (bc, snc...)

input scaling - first, it is not available in pytorch, maybe we can earase it when limiting keras support. even if not - it is like the concat threshold - a quantization enhancement

reuvenp added 30 commits June 18, 2025 12:18

build new package of graph_builder

736006c

rename refine functions and remove transforms configs

e7c1f14

create base graph builder to wrap the convertion logic

777e99a

use graph builder in ru data facade

0a03d63

add relu2pot flag in transform functions and use builder in pytorch p…

631d3fe

…tq facade

add fqc to graph and basic substitutions in transform_torch_graph

c2b6037

set prior info in nodes of torch graph

61cebd9

move torch collapsing and bn folding substitutions to new transform_g…

75d5d74

…raph

use graph builder in torch gptq facade

178eeb9

update relu_bound flag in keras builder

72c60b4

pass the builder fw class to the core runner instead of doing it outs…

8a31f5c

…ide of the core runner

use set_fqc in graph builder instead of directly fqc

227f47f

use graph builder in bn_info_collection_test

b26646e

handle tb writer in base graph builder

8114d98

fix xquant usage of graph_prepare

5922898

remove todo

fa06fc4

change visability of graph builder methods since it is used in pruning

d2ef1f6

use float graph builder in pruning torch facade

b9fa4d3

use float graph builder in pruning keras facade

11382f3

replace the usage of read model function with graph builder function

6e6ca1c

move torch reader to new package

943a844

move torch reader into new module

5a92331

adjust torch pytests

f59e6b8

remove read_model function from graph_prep_runner

bc754cf

fix keras tests

c9ab059

adjust keras pytests

ca043c0

move import of keras graph builder in ru data facade after tf check

9c95d4e

move import of torch graph builder in ru data facade after torch check

6c4175c

remove imports from graph builder init

87a1165

fix import convert_pytorch_model_to_graph

aa24e0a

reuvenp added 9 commits July 1, 2025 11:49

fix import convert_keras_model_to_graph

17adafe

fix imports for envs where torch and tf are not installed

9275ff5

remove old read_model_to_graph import in test_assert_to_operation.py

8d95ceb

adapt test ru data after removing graph_prep function with get_finali…

ca66e05

…zed_graph checks

merge with main

24e2c3c

fix functions tests after merge removal of fw graph builder

bc5a3ac

fix torch test hessian service

f85850e

fix missing import in test_weights_activation_split_substitution

fff3207

improve documentation and add licenses

86c5221

github-actions Bot added auto:core auto:gptq auto:ptq auto:qat auto:tests labels Jul 2, 2025

reuvenperetz requested a review from lior-dikstein July 2, 2025 13:26

reuvenperetz added the pr: refactoring/code cleanup label Jul 2, 2025

Merge branch 'main' into graph-prepare-refactor

289dac1

ofirgo suggested changes Jul 3, 2025

View reviewed changes

ofirgo reviewed Jul 3, 2025

View reviewed changes

reuvenperetz closed this Jul 24, 2025

		@@ -0,0 +1,14 @@
		# Copyright 2025 Sony Semiconductor Israel, Inc. All rights reserved.

		import torch


		def convert_pytorch_model_to_graph(model: torch.nn.Module,

		from model_compression_toolkit.core.common.substitutions.apply_substitutions import substitute


		def transform_pytorch_graph(graph: Graph,

		from model_compression_toolkit.graph_builder.keras.convert_keras_model_to_graph import convert_keras_model_to_graph


		class TestGraphReading(unittest.TestCase):

Conversation

reuvenperetz commented Jul 2, 2025

Pull Request Description:

Checklist before requesting a review:

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ofirgo Jul 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

ofirgo Jul 3, 2025 •

edited

Loading