Albert to torchscript is not working

[yugant-git]

Trying to export torchscript module for AlbertForQuestionAnswering.

self.model = AlbertForQuestionAnswering.from_pretrained(self.model_dir)
script_model = torch.jit.script(self.model)
script_model.save("script_model.pt")

Getting following exception:

Python builtin <built-in function next> is currently not supported in Torchscript:
	at /usr/local/lib/python3.6/dist-packages/transformers/modeling_albert.py:523:67

	device = input_ids.device if input_ids is not None else inputs_embeds.device

	if attention_mask is None:
	attention_mask = torch.ones(input_shape, device=device)
	if token_type_ids is None:
	token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)

	extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
	extended_attention_mask = extended_attention_mask.to(dtype=next(self.parameters()).dtype) # fp16 compatibility
	~~~~ <--- HERE extended_attention_mask=(1.0 - extended_attention_mask) * -10000.0 if head_mask is not None: if
		head_mask.dim()==1: head_mask=head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
		head_mask=head_mask.expand(self.config.num_hidden_layers, -1, -1, -1, -1) elif head_mask.dim()==2:
		head_mask=head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1) # We can specify head_mask for each layer
		head_mask=head_mask.to(dtype=next(self.parameters()).dtype) # switch to fload if need + fp16 compatibility
		else: '__torch__.transformers.modeling_albert.___torch_mangle_15.AlbertModel.forward' is being compiled since it
		was called from '__torch__.transformers.modeling_albert.___torch_mangle_14.AlbertForQuestionAnswering.forward'
		at /usr/local/lib/python3.6/dist-packages/transformers/modeling_albert.py:767:8 def forward(self,
		input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None,
		start_positions=None, end_positions=None): outputs=self.albert( ~~~~~~~~~~~~~~~~~~~~~~... <--- HERE
		input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids,
		head_mask=head_mask, inputs_embeds=inputs_embeds ) sequence_output=outputs[0]

[LysandreJik]

Hi, the models can be traced using the torch.jit.trace method, not the torch.jit.script. This requires inputs of the same shape that will be used for inference. Here's an example:

from transformers import AlbertForQuestionAnswering
import torch

inputs = torch.tensor([[1,2,3]])

model = AlbertForQuestionAnswering.from_pretrained("albert-base-v1")

script_model = torch.jit.trace(model, inputs)
script_model.save("script_model.pt")

[yugant-git]

thanks for the help. torch.jit.trace works. But I see that traced module perf is worse than untraced, on cpuonly mode. Any suggestion on, what I might be doing wrong.

Architecture: linux-64
OS: ubuntu-1804
GPU: None
CUDA: None
torch: 1.3.1+cpu
transformers: 2.3.0

[LysandreJik]

When tracing the model, you will need to run through it once before so that it is traced, which usually takes quite some time. This is necessary to do the just-in-time optimizations. When you run it after this, the performance should be better.

Does the performance improve after the first iteration?

[yugant-git]

My model is initialized like this.

self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(self.device)

print('loading model...')
# Load your model here
self.tokenizer = AlbertTokenizer.from_pretrained(self.model_dir)

if os.path.isfile('traced_model.pt'):
    self.model = torch.jit.load('traced_model.pt')
    print('Loading traced model')
    print(type(self.model))
else:
    self.model = AlbertForQuestionAnswering.from_pretrained(self.model_dir)
    print('Loading pytorch bin')
    print(type(self.model))

self.model.to(self.device)
self.model.eval()

The perf measurement is done like this:

model_start = time.perf_counter()
with torch.no_grad():
    if isinstance(self.model, torch.jit.ScriptModule):
        start_scores, end_scores = self.model(
            torch.tensor([all_input_ids])[0].to(self.device),
            torch.tensor([all_attention_masks])[0].to(self.device),
            torch.tensor([all_token_type_ids])[0].to(self.device)
        )
        start_scores_cpu = start_scores.cpu().tolist()
        end_scores_cpu = end_scores.cpu().tolist()
        print({ "TorchScriptExecutedInSec" : time.perf_counter() - model_start})
    else:
        start_scores, end_scores = self.model(
            torch.tensor([all_input_ids])[0].to(self.device),
            torch.tensor([all_attention_masks])[0].to(self.device),
            torch.tensor([all_token_type_ids])[0].to(self.device)
        )
        start_scores_cpu = start_scores.cpu().tolist()
        end_scores_cpu = end_scores.cpu().tolist()
        print({ "PytorchModelExecutedInSec" : time.perf_counter() - model_start})

The pytorch untraced latecy is like this (Average: 0.92225 sec):

{'PytorchModelExecutedInSec': 0.888714800003072}
{'PytorchModelExecutedInSec': 0.9285387999989325}
{'PytorchModelExecutedInSec': 0.9449487999991106}
{'PytorchModelExecutedInSec': 0.8750040000013541}
{'PytorchModelExecutedInSec': 0.9282080000011774}
{'PytorchModelExecutedInSec': 0.8841497000030358}
{'PytorchModelExecutedInSec': 0.9255469999989145}
{'PytorchModelExecutedInSec': 0.9070025000000896}
{'PytorchModelExecutedInSec': 0.9690179000026546}
{'PytorchModelExecutedInSec': 0.9713676999999734}

And traced torchscript model (Average: 0.98375664 sec):

{'TorchScriptExecutedInSec': 1.0122946000010415}
{'TorchScriptExecutedInSec': 0.9303289000017685}
{'TorchScriptExecutedInSec': 1.1499014000000898}
{'TorchScriptExecutedInSec': 1.0230705000030866}
{'TorchScriptExecutedInSec': 1.0278947000006156}
{'TorchScriptExecutedInSec': 0.9148064999972121}
{'TorchScriptExecutedInSec': 0.8976871999984724}
{'TorchScriptExecutedInSec': 0.9487294999998994}
{'TorchScriptExecutedInSec': 0.9489730000022973}
{'TorchScriptExecutedInSec': 0.9838801000005333}

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.