Search code examples
pytorchtransformer-modeltorchscript

How to use scripting to convert pytorch transformer?

I am trying to compile pytorch transformer to run it in C++:

from torch.nn import TransformerEncoder, TransformerEncoderLayer
encoder_layers = TransformerEncoderLayer(1000, 8, 512, 0.1)
transf = TransformerEncoder(encoder_layers, 6)
sm = torch.jit.script(transf)

But I am getting an error:

RuntimeError:  Expected a default value of type Tensor on parameter
"src_mask":   File "C:\Program Files (x86)\Microsoft Visual
Studio\Shared\Python36_64\lib\site-packages\torch\nn\modules\transformer.py",
line 271
    def forward(self, src, src_mask=None, src_key_padding_mask=None):
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~... <--- HERE
        r"""Pass the input through the encoder layer.

It looks like something wrong with pytorch transformer module.

Is there any way to run pytorch transformer in C++ ?

Solution

  • You need to upgrade to PyTorch 1.5.0, older versions did not support converting Transformers to TorchScript (JIT) modules.

    pip install torch===1.5.0 -f https://download.pytorch.org/whl/torch_stable.html

    In 1.5.0 you will see some warnings about the parameters being declared as constants, such as:

    UserWarning: 'q_proj_weight' was found in ScriptModule constants,  but it is a non-constant parameter. Consider removing it.

    These can be safely ignored.