[ONNX, PyTorch, TensorRT] GatherElements 지원 이슈

PyTorch 로 구현한 모델에 torch.gather 가 포함된 경우 ONNX 변환 시 ONNX 모델에  GatherElements 연산이 포함될 수 있다. 

 

onnx-tensorrt 프로젝트내에 Supported Operators에 관한 문서를 보면 GatherElements를 지원하나, TensorRT 버전에 따라 지원이 될 수도 있고 안될수도 있다. 적어도 7.2.1.6 버전에서는 지원이 안되는 것을 확인하였다. (ONNX2TensorRT 과정은 내가 맡은 부분이 아니라 확실하진 않다...)

 

https://github.com/onnx/onnx-tensorrt/blob/master/docs/operators.md

onnx/onnx-tensorrt

 

8. 대 버전부터 지원되나보다. 근데 아직(20210630 기준) 8. 대는 Early Access버전이다.

 

https://docs.nvidia.com/deeplearning/tensorrt/release-notes/tensorrt-8.html

Release Notes :: NVIDIA Deep Learning TensorRT Documentation

 

텐서알티 버전업을 하기 쉬운 환경이라면 버전업하면 해결될 문제이지만... 그게 어렵다고 하면 PyTorch2ONNX 과정에서 어떻게든 되게끔 만들어야한다... 구글링을 해보니 해결 방법이 있었다.

 

https://ask.csdn.net/questions/1480565

No importer registered for op: GatherElements-开源项目-CSDN问答

 

기존에 torch.gather를 사용한 부분에서 torch.만 없애고 아래의 함수를 거치도록 코드를 수정하면 된다.

다만, 결과는 꼭 확인해야한다. 배치 데이터를 넣었을때 모든 데이터가 제대로 처리되는지 꼭 확인해야한다.

 

 

def gather(input, dim, index):
    indices = [torch.arange(size, dtype=torch.int32, device=index.device) for size in index.shape]
    indices = list(torch.meshgrid(*indices))
    indices[dim] = index
    sizes = list(reversed(list(itertools.accumulate(reversed(input.shape), operator.mul))))
    index = sum((index * size for index, size in zip(indices, sizes[1:] + [1])))
    output = input.flatten()[index]
    return output