Another example is a feature that was just added to torchvision which can take a classification model and extract the backbone for generating embeddings.
Thanks for this reply! Got some follow up questions if you don't mind being bothered ...
> It makes it possible to automate optimizing python models, adding things like conv and batch norm fusion for inference.
By "optimize", do you mean "reduce computational load", or "use Adam/SGD/whatever to minimize a loss function"? What is "conv and batch norm fusion"? How does FX help with any of this?
> It also allows you to plug in other ops to for example make quantization or profiling easier.
I can indeed see how it could make profiling easier. I'd love to get pointers/links as to quantization methods that would necessitate adding new ops.
> Another example is a feature that was just added to torchvision which can take a classification model and extract the backbone for generating embeddings.
Hasn't it always been possible to extract a certain set of weights from some `nn.Module`?
> What is "conv and batch norm fusion"? How does FX help with any of this?
Essentially, during inference, batch norm is simply a multiply and add operation. If this occurs after a convolution, then you can simply fold (i.e. "fuse") the batch norm into the convolution by modifying the convolution's weights. See https://pytorch.org/tutorials/intermediate/fx_conv_bn_fuser.... for more details.
What the FX pass ends up looking like is:
1. Look for a convolution followed by a batch norm (where the convolution output is not used anywhere else).
2. Modify the convolution's weights to reflect the batch norm.
3. Change all users of the batch norm's outputs to use the convolution's output.
> Hasn't it always been possible to extract a certain set of weights from some `nn.Module`?
IIRC, what torchvision allows you to do now is to get a new model that simply computes the embedding.
For example, something like this (not actual code)
model = resnet18()
# Returns a new model that takes in the same inputs and returns the output at layer 3
model_backbone = create_feature_extractor(model, return_nodes=('layer3'))
Before, you'd need to manually extract out the modules that did what you wanted, and things could be tricky depending on which activation you wanted exactly. For example, if you wanted it to output one activation from inside each residual block, how would you do it?
The FX-based API allows you to simply specify what outputs you want, and it'll create a new model that provides those outputs.
It also allows you to plug in other ops to for example make quantization or profiling easier (see https://pytorch.org/tutorials/intermediate/fx_profiling_tuto...).
Another example is a feature that was just added to torchvision which can take a classification model and extract the backbone for generating embeddings.