Onnx-tensorrt: No importer registered for op: GatherElements

Created on 23 Jul 2020 · 6Comments · Source: onnx/onnx-tensorrt

When importing a model from ONNX, I get:

In node 173 (parseGraph): UNSUPPORTED_NODE: No importer registered for op: GatherElements

When will GatherElements be supported?
Or does anyone know a workaround?

Source

skmhrk1209

Most helpful comment

@bonlime Thanks. I re-implemented torch.gather using array indexing as follows:

import torch
import itertools
import operator

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

input = torch.randn(3, 3, 3)
index = torch.randint_like(input, 3, dtype=torch.long)
output1 = torch.gather(input, 1, index)
output2 = gather(input, 1, index)
assert((output1 == output2).all())

This works, and I successfully imported a model from ONNX.

skmhrk1209 on 26 Jul 2020

👍4

All 6 comments

Encountered the same problem today.
[ERROR] INVALID_ARGUMENT: getPluginCreator could not find plugin GatherElements version 1
It looks like for now the only option is to rewrite the code to avoid using it (I'm porting pytorch -> onnx first and only then onnx -> trt)

bonlime on 24 Jul 2020

FYI, I've managed to fix this problem by using usual array indexing instead of torch.gather. don't know if it could help in your case

bonlime on 26 Jul 2020

@bonlime Thanks. I re-implemented torch.gather using array indexing as follows:

import torch
import itertools
import operator

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

input = torch.randn(3, 3, 3)
index = torch.randint_like(input, 3, dtype=torch.long)
output1 = torch.gather(input, 1, index)
output2 = gather(input, 1, index)
assert((output1 == output2).all())

This works, and I successfully imported a model from ONNX.

skmhrk1209 on 26 Jul 2020

👍4

@bonlime Thanks. I re-implemented torch.gather using array indexing as follows:

import torch
import itertools
import operator

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

input = torch.randn(3, 3, 3)
index = torch.randint_like(input, 3, dtype=torch.long)
output1 = torch.gather(input, 1, index)
output2 = gather(input, 1, index)
assert((output1 == output2).all())

This works, and I successfully imported a model from ONNX.

Thank you for sharing this code, but when I implement your code it will be incorrect when input.size(dim)!=index.size(dim)
so I change it to:

import torch
import itertools
import operator

def gather(input, dim, index):
    indices = [torch.arange(size, 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

input = torch.randn(3, 3, 3)
index = torch.randint_like(input, 3, dtype=torch.long)
output1 = torch.gather(input, 1, index)
output2 = gather(input, 1, index)
assert((output1 == output2).all())

just change this line sizes = list(reversed(list(itertools.accumulate(reversed(input.shape), operator.mul))))
This will be work when dimension size is different
Thanks again!