5 ms·

What does FX enable? What are the use cases?

davidatbu 5y ago
Can someone please answer this? I'm so curious. The only real life use case that I've seen mentioned is "programmatically generating models, for example from a config file". But due to Python's dynamic nature, this is already possible. AllenNLP is a great example of that.
- m_ke 5y ago
  It makes it possible to automate optimizing python models, adding things like conv and batch norm fusion for inference. 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_tutorial.html 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.
  - davidatbu 5y ago
    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`?
    - chillee 5y ago
      > 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.html 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.
      - davidatbu 5y ago
        @chilee, you rock! I get it now.
- chillee 5y ago
  To give another take, Python's dynamic nature allows you to easily change class properties, methods, and so on. However, what is difficult with Python is to actually modify the executed code. For example, say you had def f(x): return F.relu(x) Now, you want to change all the activations in your network from a relu to gelu. This... is not so easy generically. You might be able to do it for your personal code, but what if you're importing a model from torchvision? With FX though, you can simply trace out the graph, substitute the F.relu with a F.gelu, and be done in <10 lines of code! Essentially, it gives you the freedom to perform transformations on your code (although it places limitations on what your code can contain, like no control flow).
  - Kalanos 5y ago
    Isn't that why function arguments exist? `if act=='relu'`
    - chillee 5y ago
      haha, yes, but that requires you to modify existing code to do so (which isn't always possible!). There might also be other things you want to do (like add profiling after each op) that would be tedious to do manually, but can easily automated with FX (https://pytorch.org/tutorials/intermediate/fx_profiling_tutorial.html https://pytorch.org/tutorials/intermediate/fx_profiling_tuto...). Another example is the recent support from torchvision for extracting intermediate feature activations (https://github.com/pytorch/vision/releases/tag/v0.11.0 https://github.com/pytorch/vision/releases/tag/v0.11.0). Like, sure, it was probably possible to refactor all of their code to enable users to specify extracting an intermediate feature, but it's much cleaner to do with FX.
      - davidatbu 5y ago
        Thanks @chillee! Your comments are helpful.
mirker 5y ago
My read on it is it’s equivalent to the backend of JAX’s @jit. You can write your own @jit to do operator fusion and other optimizations.
- chillee 5y ago
  Not quite - it's more of a platform for users to write their own transformations on their PyTorch code. This can include things like doing operator fusion/lowering to a backend compiler, but can also include things like inserting profiling instrumentation (https://pytorch.org/tutorials/intermediate/fx_profiling_tutorial.html https://pytorch.org/tutorials/intermediate/fx_profiling_tuto...) or extracting intermediate features (https://github.com/pytorch/vision/releases/tag/v0.11.0 https://github.com/pytorch/vision/releases/tag/v0.11.0). Basically, if you want a graph representation of a PyTorch module that's really easy to modify, use torch.fx :)
  - mirker 5y ago
    By “backend” I mean the compiler logic backing traced tensors in JAX. torch.fx seems different in that it’s primarily aimed at being a platform for users, rather than JAX which (as far as I know) hides the logic behind @jit annotations. Both trace eager code to a graph, which is then rewritten. JAX @jit is notably going eager -> graph -> XLA https://jax.readthedocs.io/en/latest/notebooks/How_JAX_primitives_work.html https://jax.readthedocs.io/en/latest/notebooks/How_JAX_primi... So (to me) it seems that they are similar backends/library primitives with different front-ends. There doesn’t seem to be a difference in representational power, since both hit a graph representation. The main exception I could see would be something like timers, which would perhaps require a graph-mode equivalent for JAX.
    - chillee 5y ago
      With this kind of stuff, I think the devil is in the details, but the principle is similar, yes. Specifically, the level of abstraction at which you're tracing, as well as what types of programs you can express/not express. For example, FX is extremely unopinionated about what it can trace, and the trace itself is extremely customizable. For example, if a subfunction/module has control flow (i.e. untraceable), it's easy to mark it as a "leaf" in FX's tracer, while that concept doesn't really make in sense in Jax's tracing system. Another example of a difference is that Jax traces out into its own IR called a jaxpr, while FX is explicitly a Python => Python translation. This has some some upsides and some downsides - for example, you can insert arbitrary Python functions into your FX graph (breakpoints, print statements, etc.), while jaxprs don't allow that. Is this a good thing? Well, if your main goal is to lower to XLA, definitely not lol. But for FX it works quite well. TL;DR: The general principles of doing graph capture are similar, but the details matter, and the details end up being quite different.

Can someone please answer this? I'm so curious. The only real life use case that I've seen mentioned is "programmatically generating models, for example from a config file". But due to Python's dynamic nature, this is already possible. AllenNLP is a great example of that.

m_ke 5y ago
It makes it possible to automate optimizing python models, adding things like conv and batch norm fusion for inference. 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_tutorial.html 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.
- davidatbu 5y ago
  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`?
  - chillee 5y ago
    > 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.html 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.
    - davidatbu 5y ago
      @chilee, you rock! I get it now.
chillee 5y ago
To give another take, Python's dynamic nature allows you to easily change class properties, methods, and so on. However, what is difficult with Python is to actually modify the executed code. For example, say you had def f(x): return F.relu(x) Now, you want to change all the activations in your network from a relu to gelu. This... is not so easy generically. You might be able to do it for your personal code, but what if you're importing a model from torchvision? With FX though, you can simply trace out the graph, substitute the F.relu with a F.gelu, and be done in <10 lines of code! Essentially, it gives you the freedom to perform transformations on your code (although it places limitations on what your code can contain, like no control flow).
- Kalanos 5y ago
  Isn't that why function arguments exist? `if act=='relu'`
  - chillee 5y ago
    haha, yes, but that requires you to modify existing code to do so (which isn't always possible!). There might also be other things you want to do (like add profiling after each op) that would be tedious to do manually, but can easily automated with FX (https://pytorch.org/tutorials/intermediate/fx_profiling_tutorial.html https://pytorch.org/tutorials/intermediate/fx_profiling_tuto...). Another example is the recent support from torchvision for extracting intermediate feature activations (https://github.com/pytorch/vision/releases/tag/v0.11.0 https://github.com/pytorch/vision/releases/tag/v0.11.0). Like, sure, it was probably possible to refactor all of their code to enable users to specify extracting an intermediate feature, but it's much cleaner to do with FX.
    - davidatbu 5y ago
      Thanks @chillee! Your comments are helpful.

mirker 5y ago

My read on it is it’s equivalent to the backend of JAX’s @jit. You can write your own @jit to do operator fusion and other optimizations.

chillee 5y ago
Not quite - it's more of a platform for users to write their own transformations on their PyTorch code. This can include things like doing operator fusion/lowering to a backend compiler, but can also include things like inserting profiling instrumentation (https://pytorch.org/tutorials/intermediate/fx_profiling_tutorial.html https://pytorch.org/tutorials/intermediate/fx_profiling_tuto...) or extracting intermediate features (https://github.com/pytorch/vision/releases/tag/v0.11.0 https://github.com/pytorch/vision/releases/tag/v0.11.0). Basically, if you want a graph representation of a PyTorch module that's really easy to modify, use torch.fx :)
- mirker 5y ago
  By “backend” I mean the compiler logic backing traced tensors in JAX. torch.fx seems different in that it’s primarily aimed at being a platform for users, rather than JAX which (as far as I know) hides the logic behind @jit annotations. Both trace eager code to a graph, which is then rewritten. JAX @jit is notably going eager -> graph -> XLA https://jax.readthedocs.io/en/latest/notebooks/How_JAX_primitives_work.html https://jax.readthedocs.io/en/latest/notebooks/How_JAX_primi... So (to me) it seems that they are similar backends/library primitives with different front-ends. There doesn’t seem to be a difference in representational power, since both hit a graph representation. The main exception I could see would be something like timers, which would perhaps require a graph-mode equivalent for JAX.
  - chillee 5y ago
    With this kind of stuff, I think the devil is in the details, but the principle is similar, yes. Specifically, the level of abstraction at which you're tracing, as well as what types of programs you can express/not express. For example, FX is extremely unopinionated about what it can trace, and the trace itself is extremely customizable. For example, if a subfunction/module has control flow (i.e. untraceable), it's easy to mark it as a "leaf" in FX's tracer, while that concept doesn't really make in sense in Jax's tracing system. Another example of a difference is that Jax traces out into its own IR called a jaxpr, while FX is explicitly a Python => Python translation. This has some some upsides and some downsides - for example, you can insert arbitrary Python functions into your FX graph (breakpoints, print statements, etc.), while jaxprs don't allow that. Is this a good thing? Well, if your main goal is to lower to XLA, definitely not lol. But for FX it works quite well. TL;DR: The general principles of doing graph capture are similar, but the details matter, and the details end up being quite different.