BERT: Cannot squeeze a dimension whose value is not 1

[jbmaxwell]

I get the following error when running prediction on pytorch-pretrained-BERT (MLM), converted from ONNX. Both ONNX and CoreML conversion complete without errors. The specific error when running in Xcode is:

[espresso] [Espresso::handle_ex_plan] exception=Espresso exception: "Invalid state": Cannot squeeze a dimension whose  
value is not 1: shape[1]=128 stat2020-02-16   
11:36:05.959261-0800 Xxxxx[6725:2140794] [coreml] Error computing NN outputs -5  

When running from python, I don't get the "Cannot squeeze" message, but do get Error computing NN outputs.

Trace (python):

Traceback (most recent call last):
  File "mlmodel_prediction_test.py", line 16, in <module>
    predictions = model.predict({'input.1': np.asarray(tokens_tensor, dtype=np.int32), 'input.3': np.asarray(segments_tensor, dtype=np.int32)})
  File "/Users/jbmaxwell/anaconda3/envs/coreml/lib/python3.7/site-packages/coremltools/models/model.py", line 334, in predict
    return self.__proxy__.predict(data, useCPUOnly)
RuntimeError: {
    NSLocalizedDescription = "Error computing NN outputs.";
}

Conversion script:

import coremltools
from onnx_coreml import convert
mlmodel = convert(model="outputs/bert_test.onnx", minimum_ios_deployment_target="13", image_input_names=["tokenizedText", "segmentsMask"], image_output_names=["predictions", "isNext"])
spec = coremltools.utils.load_spec('outputs/bert_test.mlmodel')
spec_fp16 = coremltools.utils.convert_neural_network_spec_weights_to_fp16(spec)
coremltools.utils.save_spec(spec_fp16, 'outputs/bert_test_fp16.mlmodel')
print("Saved to mlmodel (fp16): ", spec.description)

System: Train: Ubuntu 18.04 coremltools==3.3 onnx==1.6.0 onnx-coreml==1.2 python 3.7 (anaconda)

Test: macOS 10.15.3 coremltools==3.3 python 3.7 (anaconda)

[jbmaxwell]

Just as an update, I trained and converted the current huggingface distilBert model and have the same error.

[jbmaxwell]

Quick update. The above issue can be solved by modifying the mlmodel after conversion (see: https://github.com/huggingface/swift-coreml-transformers/issues/16#issuecomment-592153661). However, there seems to be a more general problem with MLM. Converting the standard Bert model completes successfully with opset v10, but DistilBert appears to require opset v11 (float input support for Equal). Though it completes, the converted standard Bert model gives incorrect output in Xcode (when compared to PyTorch). DistilBert, converted with opset v11, fails during mlmodel conversion, due to an unsupported Range operator.

If anybody has been able to get MLM working for Bert, or DistilBert, on iOS, I'd greatly appreciate any help. But since my progress seems totally blocked, I'll close this issue for now.