torch2onnx transformer encoder, decoder onnx변환시 aten::_native_multi_head_attention 에러 우회법

pytorch 의 official TransformerDecoder 레이어가 포함된 모델을 onnx 변환하려니 아래와 같이 onnx에서 지원하지 않는 연산이라며 오류가 발생했다.

 

aten::_native_multi_head_attention ~

 

우회법은 onnx 변환 함수 호출하기전에 torch.backends.mha.set_fastpath_enabled(False) 를 호출해주는 것이다.

 

이 방법 찾기 전에는 TransformerDecoder 상속 받아서 multi head attention 부분만 custom op로 돌아가게 오버라이딩 해야하나 이런 고민을 했었는데, 저거 호출하고 변환하니 일단은 변환 돼서 안심...