[ONNX] Pytorch 모델을 ONNX 모델로 변환 할 때 dynamic_axes 지정하는 방법

 

LSTM 같은 모델을 ONNX 모델로 변환할 경우, 이 모델은 입력 값을 동적으로 설정할 수 도 있기 때문에 ONNX 모델로 export 시 dynamic_axes 를 설정하여 동적인 입력 값을 갖는 ONNX 모델로 변환할 수 있다. 

 

 

torch 공식 문서에 나와있는 내용은 다음과 같다.

 

dynamic_axes (dict<string, dict<python:int, string>> or dict<string, list(int)>, default empty dict) – a dictionary to specify dynamic axes of input/output, such that:  - KEY: input and/or output names  - VALUE: index of dynamic axes for given key and potentially the name to be used for exported dynamic axes.  In general the value is defined according to one of the following ways or a combination of both:  (1). A list of integers specifying the dynamic axes of provided input. In this scenario automated names will be generated and applied to dynamic axes of provided input/output during export.  (2). An inner dictionary that specifies a mapping FROM the index of dynamic axis in corresponding input/output TO the name that is desired to be applied on such axis of such input/output during export.

 

 

 

만약에 아래와 같이 inputs 및 outputs 이 정의되어 있다면, 

 

 

 

 

dynamic_axes 은 다음과 같이 설정할 수 있다. index(indices)로 지정할 수도 있고 name 으로 지정할 수 있다. 섞어서 쓸수도 있다고 한다. 

 

 

 

 

 

LSTM 모델을 ONNX 모델로 Export 할 때 dynamic_axes 을 지정해서 변환하는 방법은 아래와 같다. 

 

 

 

 

 

참고자료 : https://pytorch.org/docs/stable/onnx.html

torch.onnx — PyTorch 1.9.0 documentation