DL| ONNX
ONNX
PyTorch를 ONNX로 export하기 (선택) PyTorch 모델을 ONNX으로 변환하고 ONNX 런타임에서 실행하기 Model Inference with ONNX
ONNX
- Open Neural Network eXchange의 줄임말
- 다른 DNN 프레임워크 환경(ex. TesorFlow, Pytorch, etc…)에서 만들어진 모델들을 서로 호환되게 사용할 수 있도록 만들어진 공유 플랫폼
장점
- Framework Interoperability: 특정 환경에서 생성된 모델을 다른 환경으로 import하여 자유롭게 사용
- Shared Optimization: 하드웨어 설계시 ONNX representation을 기준으로 최적화하면 되기 때문에 효율적
- Deep Learning Compiler(TVM, TensorRT)로 변환할 때 DNN 프레임워크를 ONNX 형태로 변환한 뒤 DL_Compiler 형태로 변환
ONNX 과정정
Pytorch -> ONNX로 export
- Pytorch 모델 & input을 인자로 torch.onnx.export 함수 호출하면 Pytorch JIT 컴파일러인 TorchScript를 통해서 trace or script 생성
- Pytorch의 nn.Module을 상속하는 모델의 forward 함수에서 실행되는 코드들에 대한 IR(Intermediate Representation)을 담고 있음
- 쉽게 말해 forward propagation 시에 호출되는 함수 및 연산들에 대한 최적화된 그래프가 만들어짐
- 생성된 trace/ script(PyTorch IR)는 ONNX exporter를 통해서 ONNX IR로 변환되고 여기서 한 번 더 graph optimization이 이뤄짐
- 최종적으로 생성된 ONNX 그래프는 .onnx 포맷으로 저장됨
TorchScript
TorchScript는 Pytorch의 Just-In-Time(JIT) 컴파일러로, PyTorch의 모델이나 함수들을 중간 표현(IR)으로 변환하여 최적화된 실행을 가능하게 함
TensorFlow는 symbolic graph execution 방식으로 작동(그래프를 먼저 정의한 후 실행) 이 과정에서 그래프 최적화가 이뤄져서 실행 성능을 높일 수 있고, 모델의 경로가 더 명확하게 최적화
TorchScript를 통해서 여러 optimization을 적용할 수 있고, serialized(직렬화)된 모델을 Pytorch dependency가 없는 다른 환경에서도 활용할 수 있는 이점이 있음
Tracing
- Pytorch에서 model(input)과 같이 nn.Module을 상속하는 모델에게 input을 넘겨주면 해당 모델 클래스의 forward 함수가 실행되면서 forward propagation이 한번 수행됨
- 이때, forward 함수 내부에서는 여러 함수가 실행이 될 수도 있는데 forward를 한번 수행하는 동안 execution path에 존재했던 모든 연산들은 IR로 기록됨
문제점
- forward에 dynamic control flow가 존재하지 않으면 trace를 생성하기 위해 한번 forward가 호출되었을 때 거쳐가지 않은 control path는 추적되지 않음 (만약 model에 if, loop와 같은 control flow가 존재하면 loop unrolling을 통해서 branch가 static하게 고정)
- 또한 trace 시에 forward 함수를 호출하기 위한 example input이 필요한데, 이 input의 shape에 따라서 graph가 고정이되어 tracing된 모델을 활용해서 traning, inferencing을 한다면 example input과 정확히 동일한 형태의 inputd을 넣어야만 한다.
전에 부스트캠프에서 어떤 것이었는지는 기억나지 않지만 ONNX 변환하는 코드를 적용했을 때 input shape이 맞지 않다는 에러가 떴었는데 이런 이유일 듯… 아!?
Script
- Trace의 단점은 Python의 dynamic feature를 살리지 못하는 것 -> 해결
- C, C++, Java와 같은 언어들에서 전체 code를 컴파일하여 사용을 하듯, scripting시에는 forward propagation 시에 실행될 전체 코드에 대해 컴파일하고 TorchScript code인 ScriptModule 인스턴스 생성
- 전체 코드를 보고 컴파일을 하기에 dynamic control flow를 살릴 수 있고, Trace의 경우처럼 example input이 필요하지도 않음
문제점
- 지원하지 않는 파이썬 코드들이 상당히 많고, type 추정, 중간에 attribute이 변하는 경우 등에 문제 발생
ONNX expert
- 변환전 추론 모드로 변경하기 위해서 torch_model.eval() 또는 torch_model.train(False)를 호출
- dropout이나 batchnorm과 같은 연산들이 추론과 학습 모드에서 다르게 작동하므로
- tracing으로 변환된 모델을 사용; 모델 변환하기 위해서 torch.onnx.export() 함수를 호출
ONNXRuntime is faster than Pytorch
ONNX로 변환 후 ONNXRuntime을 사용하여 추론엔진을 구축하는 방법이 있음 ONNXRuntime은 Pytorch만으로 추론했을 때 속도가 향상
왜냐면 ONNXRuntime graph는 static graph 방식을 사용하기에 준비 과정에서 모델의 모든 구조를 보고 모델과 하드웨어에 맞게 최적화한 후 엔진에 로딩하기에 효율적이게 추론
부스트캠프 파이널 프젝에서 사용한 코드.. ONNX를 사용했을 때 50%정도 ms가 빨랐던 것으로 기억
import onnxruntime as ort
def init_onnx_sessions():
dummy_image = np.random.randint(0, 256, (640, 640, 3), dtype=np.uint8)
providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
global session_yw, session_da
session_yw = ort.InferenceSession(YOLO_ONNX_PATH, providers=providers)
session_da = ort.InferenceSession(DA_ONNX_PATH, providers=providers)
inference_yw(dummy_image)
inference_da(dummy_image)
