Optimization
Last updated
Last updated
The optimum.fx.optimization
module provides a set of torch.fx graph transformations, along with classes and functions to write your own transformations and compose them.
In 🌍 Optimum, there are two kinds of transformations: reversible and non-reversible transformations.
The most basic case of transformations is non-reversible transformations. Those transformations cannot be reversed, meaning that after applying them to a graph module, there is no way to get the original model back. To implement such transformations in 🌍 Optimum, it is very easy: you just need to subclass and implement the method.
For instance, the following transformation changes all the multiplications to additions:
Copied
After implementing it, your transformation can be used as a regular function:
Copied
For instance, the following transformation is reversible:
Copied
Copied
A reversible transformation implements both the transformation and its reverse, allowing to retrieve the original model from the transformed one. To implement such transformation, you need to subclass and implement the and methods.
As applying multiple transformations in chain is needed more often that not, is provided. It is an utility function that allows you to create a transformation by chaining multiple other transformations.