Skip to content

Draft: Rewrite memory manager log_info in Python.

Cyril Moineau requested to merge meminfo_plot into dev

Context

This MR fix !276

It began by looking at the memory management of the dinov2 model

memory_info_old

Realizing that it is not readable, I rewrited it in python in order to easily tweak it.

Here is the new memory info after rewritting it:

memory_info_name

I added an option display_names whcih allow to diplay the names or not:

memory_info_no_name

Benefice

This MR remove a dependencie to gnuplot 😄

Log function is easy to use now!

Miscellaneous

After using MetaOperators to simplify the graph we get the follwing memory placement strategy which surprised me as it produces a lower memory peak!

This make sens since by fusing operation we load in RAM data and do not need to save intermediate output.

A little sad to not have found this out before webinar 😞

memory_info_simplified

For reproductibility the code used:

import aidge_core
import aidge_backend_cpu
import aidge_onnx

dinov2_model = aidge_onnx.load_onnx("hf_dinov2_sim.onnx")

aidge_core.fuse_to_metaops(dinov2_model, "MatMul-*>Add", "Linear")

aidge_core.fuse_to_metaops(dinov2_model, "ReduceMean-*>Sub#1~>(Pow#1->ReduceMean-*>Add#1->Sqrt)-*>Div#1-*>Mul#1-*>Add#2;"
                                        "Sub#1~*>Div#1;"
                                        "Pow#1<1~Producer;"
                                        "Add#1<*~Producer;"
                                        "Mul#1<*~Producer;"
                                        "Add#2<*~Producer;"
                                        "Sub#1~>$", "LayerNorm")

aidge_core.fuse_to_metaops(dinov2_model, "MatMul->Div#1->Softmax-*>MatMul;"
                                        "Div#1<1~Producer", "ScaledDotProductAttention")

aidge_core.fuse_to_metaops(dinov2_model, "ScaledDotProductAttention#1->Transpose->Reshape#1->Linear;"
                                        "Reshape#1<1~Producer;"
                                        "ScaledDotProductAttention#1<0-(Transpose<-Reshape#2<-Add#1);"
                                        "ScaledDotProductAttention#1<1-(Transpose<-Reshape#3<-Add#2);"
                                        "ScaledDotProductAttention#1<2-(Transpose<-Reshape#4<-Add#3);"
                                        "Reshape#2<1~Producer;"
                                        "Add#1<*-0-Split#1;"
                                        "Add#2<*-1-Split#1;"
                                        "Add#3<*-2-Split#1;"
                                        "Split#1<-MatMul;"
                                        "Split#1<1~Producer", "MultiHeadAttention")

aidge_core.fuse_to_metaops(dinov2_model, "Div#1->Erf->Add#1-*>Mul->Mul#2;"
                                        "Div#1<1~Producer;"
                                        "Add#1<*~Producer;"
                                        "Mul#2<*~Producer", "GeLU")
dinov2_model.set_ordered_outputs([dinov2_model.get_ordered_outputs()[0][0].inputs()[0], dinov2_model.get_ordered_outputs()[0]])
dinov2_model.set_backend("cpu")
dinov2_model.set_datatype(aidge_core.dtype.float32)
dinov2_model.forward_dims([[1,3,224,224]], True)

s = aidge_core.SequentialScheduler(dinov2_model)
s.generate_scheduling()
aidge_core.generate_optimized_memory_info(s, "stats_dyno", wrapping=False, display_names=False)
Edited by Cyril Moineau

Merge request reports

Loading