polodealvarado
commented
1 year ago Now running onnx_implementation.py works.
The ONNX validation is done with a tolerance error of less than 0.0001.
================ Diagnostic Run torch.onnx.export version 2.0.1 ================
verbose: False, log level: Level.ERROR
======================= 0 NONE 0 NOTE 0 WARNING 0 ERROR ========================
Optimum version: 1.13.2
Validating ONNX model spanmarker_model.onnx...
-[✓] ONNX model output names match reference model (last_hidden_state, pooler_output)
- Validating ONNX Model output "last_hidden_state":
-[✓] (2, 512, 768) matches (2, 512, 768)
-[✓] all values close (atol: 0.0001)
- Validating ONNX Model output "pooler_output":
-[✓] (2, 768) matches (2, 768)
-[✓] all values close (atol: 0.0001)I have splitted the forward function into two parts: forward and _postprocessing_encoderoutput:
def forward(
self,
input_ids: torch.Tensor,
attention_mask: torch.Tensor,
position_ids: torch.Tensor,
**kwargs,
) -> Dict[str, torch.Tensor]:
"""Forward call of the SpanMarkerModel.
Args:
input_ids (~torch.Tensor): Input IDs including start/end markers.
attention_mask (~torch.Tensor): Attention mask matrix including one-directional attention for markers.
position_ids (~torch.Tensor): Position IDs including start/end markers.
None.
Returns:
outputs: Encoder outputs
"""
token_type_ids = torch.zeros_like(input_ids)
outputs = self.encoder(
input_ids,
attention_mask=attention_mask,
token_type_ids=token_type_ids,
position_ids=position_ids,
)
return outputs
def postprocessing_encoder_output(
self,
encoder_outputs: Dict[str, torch.Tensor],
start_marker_indices: torch.Tensor,
num_marker_pairs: torch.Tensor,
labels: Optional[torch.Tensor] = None,
num_words: Optional[torch.Tensor] = None,
document_ids: Optional[torch.Tensor] = None,
sentence_ids: Optional[torch.Tensor] = None,
**kwargs,
) -> SpanMarkerOutput:
"""
Post-processes the output from the encoder to produce SpanMarkerOutput.
Args:
encoder_outputs (Dict[str, torch.Tensor]): Outputs from the encoder, typically including the last hidden state.
start_marker_indices (torch.Tensor): Indices in the input sequence where start markers are located.
num_marker_pairs (torch.Tensor): Number of start-end marker pairs in each example in the batch.
labels (Optional[torch.Tensor]): Labels for the input data, used for supervised learning. Default is None.
num_words (Optional[torch.Tensor]): Number of words in each input sequence. Default is None.
document_ids (Optional[torch.Tensor]): Identifiers for the documents in the batch. Default is None.
sentence_ids (Optional[torch.Tensor]): Identifiers for the sentences in the batch. Default is None.
Returns:
SpanMarkerOutput: Processed features and optionally the loss, if labels are provided.
"""
last_hidden_state = encoder_outputs[0]
last_hidden_state = self.dropout(last_hidden_state)
batch_size = last_hidden_state.size(0)
sequence_length = last_hidden_state.size(1)
# Get the indices where the end markers start
end_marker_indices = start_marker_indices + num_marker_pairs
# The start marker embeddings concatenated with the end marker embeddings.
# This is kind of breaking the cardinal rule of GPU-based ML, as this is processing
# the batch iteratively per sample, but every sample produces a different shape matrix
# and this is the most convenient way to recombine them into a matrix.
embeddings = []
for i in range(batch_size):
embeddings.append(
torch.cat(
(
last_hidden_state[i, start_marker_indices[i] : end_marker_indices[i]],
last_hidden_state[i, end_marker_indices[i] : end_marker_indices[i] + num_marker_pairs[i]],
),
dim=-1,
)
)
padded_embeddings = [
F.pad(embedding, (0, 0, 0, sequence_length // 2 - embedding.shape[0])) for embedding in embeddings
]
feature_vector = torch.stack(padded_embeddings)
# NOTE: This was wrong in the older tests
feature_vector = self.dropout(feature_vector)
logits = self.classifier(feature_vector)
if labels is not None:
loss = self.loss_func(logits.view(-1, self.config.num_labels), labels.view(-1))
return SpanMarkerOutput(
loss=loss if labels is not None else None,
logits=logits,
*encoder_outputs[2:],
num_marker_pairs=num_marker_pairs,
num_words=num_words,
document_ids=document_ids,
sentence_ids=sentence_ids,
)
tomaarsen
Hello Tom!
Here you have the first draft for the onnx implementation.
I propose the following to-do list to get the ONNX format for SpanMarkerModel mainly based on the Optimum documentation of HuggingFace :
Requirements
How to get the onnx model
If you run the script _onnximplementation.py you get the model exported into onnx format. The model is perfectly generated but we should fix a problem with the input/output variable names.
Error response
There is a problem with the names of the generated input/output nodes because ONNX sets different names for each node. In the forward function the following input names are the same in the outputs when the SpanMarkerOutput object is created: "start_marker_indices", "num_marker_pairs", "num_words", "document_ids", "sentence_ids". By default, ONNX transforms the corresponding input names by adding '.1' at the end (You can check the error response at the bottom of the script).
Reviewing the code, I noticed that within the 'forward' function there are 'post-processing' stages included after the encoder output is obtained. One thing we could do is to split this code and leave the 'forward' solely for the execution of the encoder and then apply the post processing steps and the input/output names are not affected. This way the input/output names are unique.