Skip to content

Commit

Permalink
fix torchscript docs (huggingface#6740)
Browse files Browse the repository at this point in the history
  • Loading branch information
patrickvonplaten authored and fabiocapsouza committed Nov 15, 2020
1 parent eaeea95 commit 392d8aa
Showing 1 changed file with 9 additions and 10 deletions.
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

0 comments on commit 392d8aa

Please sign in to comment.