Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Torchscript] Fix docs #6740

Merged
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 9 additions & 10 deletions docs/source/serialization.rst
Original file line number Diff line number Diff line change
Expand Up @@ -130,13 +130,12 @@ Pytorch's two modules `JIT and TRACE <https://pytorch.org/docs/stable/jit.html>`
their model to be re-used in other programs, such as efficiency-oriented C++ programs.

We have provided an interface that allows the export of 🤗 Transformers models to TorchScript so that they can
be reused in a different environment than a Pytorch-based python program. Here we explain how to use our models so that
they can be exported, and what to be mindful of when using these models with TorchScript.
be reused in a different environment than a Pytorch-based python program. Here we explain how to export and use our models using TorchScript.

Exporting a model needs two things:
Exporting a model requires two things:

* dummy inputs to execute a model forward pass.
* the model needs to be instantiated with the ``torchscript`` flag.
* a forward pass with dummy inputs.
* model instantiation with the ``torchscript`` flag.

These necessities imply several things developers should be careful about. These are detailed below.

Expand All @@ -147,8 +146,8 @@ Implications
TorchScript flag and tied weights
------------------------------------------------
This flag is necessary because most of the language models in this repository have tied weights between their
``Embedding`` layer and their ``Decoding`` layer. TorchScript does not allow the export of models that have tied weights,
it is therefore necessary to untie the weights beforehand.
``Embedding`` layer and their ``Decoding`` layer. TorchScript does not allow the export of models that have tied weights, therefore
it is necessary to untie and clone the weights beforehand.

This implies that models instantiated with the ``torchscript`` flag have their ``Embedding`` layer and ``Decoding`` layer
separate, which means that they should not be trained down the line. Training would de-synchronize the two layers,
Expand Down Expand Up @@ -181,7 +180,7 @@ when exporting varying sequence-length models.
Using TorchScript in Python
-------------------------------------------------

Below are examples of using the Python to save, load models as well as how to use the trace for inference.
Below is an example, showing how to save, load models as well as how to use the trace for inference.

Saving a model
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Expand Down Expand Up @@ -237,10 +236,10 @@ We are re-using the previously initialised ``dummy_input``.

.. code-block:: python

loaded_model = torch.jit.load("traced_model.pt")
loaded_model = torch.jit.load("traced_bert.pt")
loaded_model.eval()

all_encoder_layers, pooled_output = loaded_model(dummy_input)
all_encoder_layers, pooled_output = loaded_model(*dummy_input)

Using a traced model for inference
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Expand Down