Repository navigation
docs: document the Torch backend and checkpointing prepared models - #175
Conversation
Add two Introduction pages: - Fast inference in PyTorch: what ExportBackend.Torch does, and when to use it and when not to. - Checkpoints and workflows: how to save and reload a prepared model with prepare(..., state_dict=...), with workflows for evaluating with the Torch backend then exporting to Core AI, resuming pruning with fine-tuning, and checkpointing a P4A8 model. Link them from "How to use coreai-opt" and the quantization, palettization and pruning overviews, and show the Torch backend next to the CoreAI and CoreML finalize examples.
| The finalized model can then be used to convert with [coreai-torch](https://github.com/apple/coreai-torch) to produce a Core AI model (`.aimodel`). | ||
| For details, see [Integration with Core AI](integration_coreai.md). | ||
| `finalize()` takes a `backend` argument. With `ExportBackend.CoreAI` (the default), the finalized model can be converted with [coreai-torch](https://github.com/apple/coreai-torch) to produce a Core AI model (`.aimodel`); see [Integration with Core AI](integration_coreai.md). `ExportBackend.CoreML` prepares the model for conversion with `coremltools`, while `ExportBackend.Torch` folds the compression into the weights for fast inference in PyTorch, for example to evaluate the compressed model; see [Fast inference in PyTorch](torch_backend.md). | ||
|
|
There was a problem hiding this comment.
"folds the compression into the weights for fast inference in PyTorch" ---> I think we should probably be more precise here. So sth like "fold (i.e. de-compress) the compressed weights, so that the resulting torch model will run at similar latency as the uncompressed weights model"
| @@ -0,0 +1,24 @@ | |||
| # Fast inference in PyTorch | |||
There was a problem hiding this comment.
Similar comment as this applies to this page.
-
the title: instead of saying "fast inference in torch" (which assumes something is slow, but its not clear what the context is), it should probably say "Running Torch evaluations with compressed models"
-
Instead of starting with "A prepared model applies compression...." , more context and setup should be provided. e.g. "After a model is compressed with coreai-opt, using the prepare API (link to it) and (optionally) the calibration/training steps, the model is ready to be evaluated and/or exported to Core AI/Core ML. The
finalizeAPI helps in doing this. There are 3 backends:...." then go on to describe.
| @@ -0,0 +1,186 @@ | |||
| # Checkpoints and workflows | |||
|
|
|||
| A prepared model can be saved as a checkpoint and reloaded later, with its compression state intact. For example, you can use a checkpoint to: | |||
There was a problem hiding this comment.
On any new page, when a "prepared" or "finalized" model is referred to, dont assume people know what it is. Pls add the hyperlink to the the appropriate section in "How to use coreai-opt" page.
| prepared_model(batch) | ||
| torch.save(prepared_model.state_dict(), "quantized.pt") | ||
|
|
||
| # Later: rebuild the P4A8 model from the checkpoints. |
There was a problem hiding this comment.
In this code snippet its not clear why one would need to "rebuild the P4A8 model".
Is it because one would first want to call finalize with torch backend, do evals, and the do the model export? If yes, then we should add the eval step also in the code
|
|
||
| This page covers how to checkpoint and reload a prepared model, and shows workflows that use checkpoints. | ||
|
|
||
| ## Checkpointing and reloading prepared models |
There was a problem hiding this comment.
It would be good to have a section on checkpointing/reloading finalized models as well. If that works out of the box like with regular torch models, it would be good to say that explicitly as well
There was a problem hiding this comment.
I agree, it is a good idea. But before documenting it, we need to check what works for each backend combination. I'll file a separate issue for that rather than expand this PR. For now, the recommended path is to checkpoint the prepared model and finalize it again, as described on this page.
| @@ -0,0 +1,30 @@ | |||
| # Running PyTorch evaluations with compressed models | |||
There was a problem hiding this comment.
Do we want to frame .finalize(ExportBackend.Torch) as only for running PyTorch evaluations?
I think we can put this as a part of the "integration with CoreAI" which is also where we talk about CoreML finalize as well or we make a new section on "finalize options" and talk about each of the finalize options individually.
There was a problem hiding this comment.
I prefer keeping it as a separate page so we can expand it later with optimizations such as using CUDA or Triton kernels. It is mainly for PyTorch evaluations and distinct from exporting for on-device use, as is the case with CoreAI or CoreML.
- Retitle the Torch backend page and introduce prepare(), finalize() and the three backends before describing the Torch backend. - Describe the Torch backend as running at close to the uncompressed model's latency instead of calling it fast. - Link prepared and finalized models to How to use coreai-opt. - Explain when to rebuild a P4A8 model from checkpoints.
7e0e479 to
f1d129b
Compare
Summary
Documents the
ExportBackend.Torchbackend (#154) and checkpointing prepared models withprepare(..., state_dict=...)(#157).Changes
how_to_use_coreaiopt.mdand the quantization, palettization and pruning overviews;Torchadded next to the CoreAI/CoreMLfinalizeexamples.