[ONNX] Pytorch 에서 Onnx 로 변환 시 Gather op 때문에 export 안되는 문제

728x90

반응형

Pytorch 에서 Onnx 모델로 변환시 Gather 와 같은 옵션 때문에 변환이 안되는 문제가 발생한다. 

이유는 Onnx 에서 Gather 라는 옵션을 지원하지 않기 때문이다. (2020.04.28 기준)

 

아래와 같이 interpolate 를 scale factor 로 구현하였을 때 

class ResizeModel(nn.Module):
    def __init__(self):
        super(ResizeModel, self).__init__()
    def forward(self, x):
        return F.interpolate(x, scale_factor=(2, 2), mode='nearest')

생성되는 그래프는 아래와 같다

 

 

 

위와 같이 구현할 경우 "Gather" 라는 op 이 붙어서 그래프가 생성되는데, 

Onnx 에서는 Gather 라는 옵션을 지원하지 않기 때문에 (2020.04.28 기준)

아래와 같이 interpolate의 scale factor 를 구현해야한다.

class ResizeModel(nn.Module):
    def __init__(self):
        super(ResizeModel, self).__init__()
    def forward(self, x):
        sh = torch.tensor(x.shape)
        return F.interpolate(x, size=(sh[2] * 2, sh[3] * 2), mode='nearest')

그러면 Gather 옵션을 피해서 아래와 같이 그래프가 생성된다.

 

 

 

 

 

 

 

 

참고자료 : https://github.com/onnx/onnx-tensorrt/issues/192

Gather in Upsample problem · Issue #192 · onnx/onnx-tensorrt

 

 

728x90

반응형

저작자표시비영리

• 카카오스토리
• 트위터
• 페이스북

'AI Development > ONNX' 카테고리의 다른 글

[ONNX] ONNX 배치 사이즈 변경하는 방법 + 삽질  (3)	2021.02.02
[ONNX] Pytorch 모델을 ONNX 모델로 변환하기  (4)	2020.08.23
[ONNX] onnx-graphsurgeon 이용하여 plugin 사용하기 - Group Normalization  (3)	2020.07.21
[ONNX] Netron : ONNX model Visualization  (0)	2020.07.16
[Onnx] Onnx Tutorials (정리중)  (0)	2020.05.14
[ONNX] Pytorch 에서 Onnx 로 변환 시 Gather op 때문에 export 안되는 문제  (5)	2020.04.28