compile a model using scripting

[joaomsimoes]

Hi everyone,

Has anyone tried to compile the model using TorchScript?

If I try:

model = whisper.load_model("small.en")
scripted_model = torch.jit.script(model)

Gives me:

RuntimeError: 'Tensor' object has no attribute or method '_conv_forward'.: File "E:\hodl-out\nlp_venv\lib\site-packages\whisper\model.py", line 43 def _conv_forward(self, x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> Tensor: return super()._conv_forward( ~~~~~~~~~~~~~~~~~~~ <--- HERE x, weight.to(x.dtype), None if bias is None else bias.to(x.dtype) ) 'Conv1d._conv_forward' is being compiled since it was called from 'Conv1d.forward' File "E:\hodl-out\nlp_venv\lib\site-packages\torch\nn\modules\conv.py", line 313 def forward(self, input: Tensor) -> Tensor: return self._conv_forward(input, self.weight, self.bias) ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE

Thanks in advance

[glangford]

Haven't tried it João, but could this have something to do with it:

If you’re using Sequential with TorchScript, the inputs of some of the Sequential submodules may be falsely inferred to be Tensor, even if they’re annotated otherwise. The canonical solution is to subclass nn.Sequential and redeclare forward with the input typed correctly.

From the Known Issues section of
https://pytorch.org/docs/stable/jit.html

[ch3njust1n]

@joaomsimoes were you able to figure this out?

[magicse]

@joaomsimoes
Try separate encoder and decoder.... something like this

import os
import torch
import torchvision.models as models
import whisper

name = "medium"

net = whisper.load_model(name, download_root=os.path.join(os.getcwd(), ""))
net = net.encoder.eval()

x = torch.rand(1, 80, 3000)
mod = torch.jit.trace(net, x)

mod.save("medium_encoder_jit_trace.pt")