Package: torch-dag
1. What istorch-dag?torch-dagis a repository in which we implement a graph-like representation fortorchmodels so that we have a
unified structure for every model. This structure is a directed acyclic graph (DAG), which we call aDagModule. We do
this so that:We can easily run graph modification algorithms on the whole model without the need to edit any code defining the
model. For example:we can easily switch all activation functions in the model with just a couple lines of codewe can easily insert new vertices (modules) into the model based on some predefined pattern, e.g., add batch norm
layers after every convolution provided there is none present immediately after the convolutionwe can introduce hierarchy into the model by wrapping some parts of the graph -> this is very useful if one wants to
run some Neural Architecture Search algorithmsWe can build a model once in the plainpytorchway, then convert it to aDagModule, save it and load itwithout
having to have the source code for model building. This is extremely useful - imagine all the trouble one has to go
through when one needs to modify the model source code manually to some simple stuff like changing some of the model
layers.The conversion between a regulartorch.nn.Moduleinstance andDagModuleusestorch.FX. We decide not to
usetorch.FXgraph representation directly for the following reasons:there is no serialization/deserialization
fortorch.FXGraphModulethe graph generated bytorch.FXis often not a computational deep learning graph one would like to work with since
it contains a huge number of purelypythonoperations as nodesthere is no notion of hierarchy/nesting inGraphModulethere is often a number of ways to represent the same operation intorch.FXGraphModule(think of addition
asx+y,torch.add(x, y), which would have a very different representation inGraphModule), which leads to plenty
of issues when one tries to implement some method to modify graph based on the presence of additions -> one has to
think of all the different ways in which addition can be implemented intorch!Note: Not everytorch.nn.Modulemodel can be converted to aDagModule! We mostly share the limitations
oftorch.FXin this respect, although we try to overcome that by explicitly considering sometorch.FX-unfriendly
modules as atomic modules of our representation.2. Why do we need a unified DAG-like representation?The short answer isto run algorithmsfor stuff like model compression. In torch-dag we implement a channel pruning
algorithm that aspires to be architecture agnostic and we try really hard to make it work for as many architectures as
possible. The exhaustive list of thetimmmodels we support is givenhere.3. Installation3.1 Installation from cloned source codeIn the cloned project's root (where this README.md file is) runpipinstall-e.3.2 Installation from PyPIpipinstalltorch-dag4. BasicsIf you have a atoch.nn.Modulemodel and you want to convert it to aDagModulesimply run:importtorch_dagastddag_model=td.build_from_unstructured_module(model)For details and more extended documentation seeHow to convert torch.nn.Module instances to DagModule?5. Model compression algorithms5.1 Channel-pruningAjupyternotebook with atoy example of channel pruning can be viewedhere.
If you want to read a more detailed intro to channel pruning withtorch-daghavce a look atpruning readme.This is the algorithm we spent plenty of time developing and refining. It helped us toWINMobile AI 2022 Single-Image Depth Estimation on Mobile Devices.Supported Models and LimitationsCurrent model coveragetimm==0.9.5num modelspercentageall models945100.00%convertible models69073.02%channel prunable_models58561.90%At this point we support plenty of convolutional models and a subset of vision transformer architectures.
A full list of supportedtimmmodels and a proportion of FLOPS that can be removed in each mdoel can be seenchannel pruning supported models.Wedo NOTsupport models that cannot be traced usingtorch.FX(there are notable exceptions, that require
additional processing likevit*anddeit3*models, whicharesupported)
seelimitations-of-symbolic-tracing. If you
have a custom model that can be traced:my_module = Model()
from torch.fx import symbolic_trace
# Symbolic tracing frontend - captures the semantics of the module
symbolic_traced : torch.fx.GraphModule = symbolic_trace(my_module)then chances are it will be convertible to aDagModuleand some part of it can be pruned! How much precisely?Results inImageNet1kfor sometimmmodels.NOTE: as a proxy for model size we often useFLOPsnormalized by input resolution expressed in thousands
(kmappfor short - thousands ofFLOPsper pixel), i.e.:kmapp(model) = FLOPs(model) / (H * W * 1000), where(H, W)is
input resolution.fbnetv3_g.ra2_in1kFBNETV3is a highly optimized architecture so further slimming down
of it is challenging. Nonetheless, the results are still pretty good. We prunefbnetv3gand then fine-tune the pruned
model for 200 epochs.Results onImageNet1Kvalidation set.modelacc (224x224)acc (240x240)GFLOPs (224x224)params (m)FLOPs reductionfbnetv3g0.80610.81322.1416.6fbnetv3d0.78560.79271.0410.3fbnetv3b0.78120.78710.8458.6m160.77420.77930.7997.8↓ 62.5%m220.79200.79551.210.5↓ 43.9%m280.796860.80161.39611.6↓ 34.8%m320.8030.80251.6113.1↓ 24.8%convnextv2_nano.fcmae_ft_in22k_in1kConvNextV2is a purely convolutional model that is inspired
by the design of Vision Transformers. We pruneconvnextv2_nano.fcmae_ft_in22k_in1kand then fine-tune the pruned
model for 200 epochs.Results onImageNet1Kvalidation set.modelacc (224x224)GFLOPsparams (m)FLOPs reductionbaseline0.81974.9115.6m0.50.81552.649.4↓  46.2%m0.30.79221.686.3↓  65.8%m0.20.75310.933.8↓ 81.0%To see more much more results have a lookhereTODO Block-pruningTODO Learnable Low Rank Compression (LLRC)