search

huggingface / swift-coreml-transformers

Swift Core ML 3 implementations of GPT-2, DistilGPT-2, BERT, and DistilBERT for Question answering. Other Transformers coming soon!
Apache License 2.0
1.61k stars 173 forks source link

distilbert-onnx-coreml.py Segmentation fault attempting MLM version #17

Closed jbmaxwell closed 4 years ago

jbmaxwell commented 4 years ago

I've trained a version of DistilBertForMaskedLM and am trying to convert it to CoreML. The only way I adapted the script, so far, was to remove the custom_conversion_functions={"Softmax": _convert_softmax} option in the convert() function, since I just want the predictions out, but it exits with a Segmentation Fault. The last few lines of the trace appear to just indicate the completion of the onnx export: 

%1107 : Float(1, 256, 512) = onnx::Add(%1106, %vocab_layer_norm.bias) # /home/james/anaconda3/envs/torch/lib/python3.7/site-packages/torch/nn/functional.py:1696:0
  %1108 : Float(512, 219) = onnx::Transpose[perm=[1, 0]](%distilbert.embeddings.word_embeddings.weight) # /home/james/anaconda3/envs/torch/lib/python3.7/site-packages/torch/nn/functional.py:1372:0
  %1109 : Float(1, 256, 219) = onnx::MatMul(%1107, %1108) # /home/james/anaconda3/envs/torch/lib/python3.7/site-packages/torch/nn/functional.py:1372:0
  %output_weights : Float(1, 256, 219) = onnx::Add(%1109, %vocab_projector.bias) # /home/james/anaconda3/envs/torch/lib/python3.7/site-packages/torch/nn/functional.py:1374:0
  return (%output_weights)

Segmentation fault (core dumped)

I'm guessing that maybe the final layers of DistilBertForMaskedLM need custom functions (vocab_transform, vocab_layer_norm, and vocab_projector)?

Any help/tips greatly appreciated.

UPDATE: Looking into DistilBertForMaskedLM, I decided to try just exporting the underlying model using model = model.distilbert, but I get the same Segmentation Fault.

UPDATE 2: Oh, gulp... Wait, I'm running this on my Ubuntu machine (where I do the training). Probably the issue... (Confirmed.)