MLOps study - Raviraja Week 4: ONNX
22 Oct 2022Week 4는 ONNX에 대한 내용입니다. 우리는 일반적으로 PyTorch나 TensorFlow를 통해 모델을 학습합니다. 그렇지만 inference는 어디에서 진행될까요? 사용하는 사람마다 다르겠죠. PC일 수도 있고 mobile 환경일 수도 있습니다. 학습은 PyTorch 이지만 inference는 TensorFlow에서 수행해야 할지도 모릅니다. 이러한 상황을 위해서 ONNX를 이용한 Model packaging에 대해서 알아보도록 하겠습니다.
Start ONNX
학습된 model을 ONNX 모델로 만들어준 후, inference 할 때 사용자가 원하는 framework를 사용할 수 있도록 만들어주는 과정을 model packaging 이라고 합니다.
PyTorch, PyTorch-lightning 모델에 대해서 전부 간단하게 동작합니다만 PyTorch-lightning의 경우 모델 자체에 onnx로 변환하는 메쏘드가 포함되어 있습니다. onnx 모델을 만드는 데 필요한 정보는 아래와 같습니다.
- Name of the onnx model
- Input sample
- Input names (초기 input 뿐만 아니라 각 layer의 input들까지 이름을 지정해 줄 수 있다. layer input의 개수보다 적을 시에는 남는 것은 자동으로 이름 붙여진다.)
- Output names (output의 이름)
- Dynamic axes (각 input들에서 변하도록 설계된 axes를 표시해 주는 것. 일반적으로는 batch인 0번 axis 혹은 RNN의 sequence length 정도가 되겠다.)
아래 코드를 돌려주면 .onnx 파일이 생성됩니다!
#### Model과 Sample 준비 과정
model_path = f"models/best-checkpoint.ckpt"
cola_model = ColaModel.load_from_checkpoint(model_path)
data_model = DataModule()
data_model.prepare_data()
data_model.setup()
input_batch = next(iter(data_model.train_dataloader()))
input_sample = {
"input_ids": input_batch["input_ids"][0].unsqueeze(0),
"attention_mask": input_batch["attention_mask"][0].unsqueeze(0),
}
#### When using PyTorch
torch.onnx.export(
cola_model, # model being run
(
input_sample["input_ids"],
input_sample["attention_mask"],
), # model input (or a tuple for multiple inputs)
"models/model.onnx", # where to save the model
export_params=True,
opset_version=10,
verbose=True,
input_names=["input_ids", "attention_mask"], # the model's input names
output_names=["output"], # the model's output names
dynamic_axes={ # variable length axes
"input_ids": {0: "batch_size"},
"attention_mask": {0: "batch_size"},
"output": {0: "batch_size"},
},
)
##### When using PyTorch-lightning
cola_model.to_onnx(
"models/model.onnx", # where to save the model
input_sample, # input samples with atleast batch size as 1
export_params=True,
opset_version=10,
input_names = ["input_ids", "attention_mask"], # Input names
output_names = ['output'], # Output names
dynamic_axes={ # variable length axes
"input_ids": {0: "batch_size"},
"attention_mask": {0: "batch_size"},
"output": {0: "batch_size"},
},
)
ONNX runtime
ONNX runtime은 ONNX 모델의 inference engine 입니다. 우선 Cuda 버전에 맞는 버전을 찾아 설치를 해줍시다. Cuda 버전에 맞지 않으면 GPU를 사용하는 inference를 할 수 없습니다. (Cuda-ONNXruntime version table)
python -m pip install onnxruntime-gpu==<version>
ONNX runtime은 서로 다른 OS와 HW(accelerator)에서 간편히 동작할 수 있도록 만들어졌습니다. HW라는 것은 GPU나 NPU, TPU 이런 것들을 의미합니다. 또한 Python 뿐 아니라 C++, JAVA, Ruby를 비롯한 다양한 언어를 사용하여 라이브러리가 구성되어 있다고 공식 홈페이지에 설명되어 있네요.
가능한 모든 HW와 현재 플랫폼에서 사용 가능한 HW는 아래와 같이 확인 가능합니다. 안타깝게도 제 연구실 서버는 Cuda 9.1이라서 onnxruntime-gpu를 지원하지는 않는 것 같습니다…ㅠㅠ
from onnxruntime import get_all_providers, get_available_providers
print(get_all_providers())
print(get_available_providers())
이걸 이용해서 inference를 한 번 해보겠습니다. onnx 모델을 불러와 InferenceSession으로 만들고 input과 함께 run 만 해주면 됩니다. output_names 의 경우 특정 이름의 output만을 결과로 얻고 싶을 때 사용합니다. None으로 지정하면 모든 output을 return 합니다.
import onnxruntime as ort
onnx_model_path = 'models/model.onnx'
ort_session = ort.InferenceSession(onnx_model_path)
ort_inputs = {
"input_ids": input_sample["input_ids"].numpy(),
"attention_mask": input_sample["attention_mask"].numpy(),
}
output_name = None
ort_output = ort_session.run(output_name, ort_inputs)
재미있는 건 ONNX 가 PyTorch inference보다 훨씬 빠르다는 겁니다 (2~3배 정도). 그런데 구글에선 아무도 왜 빠른지에 의문을 품지는 않고 있네요…(그냥 최적화를 잘 한걸지도…??) 사실 Pytorch는 꽤나 불필요한 연산들이 많이 진행될 텐데요, gradient도 계산한 것을 전부 저장하고 있고 그러니까요. 그런거 전부 날리면 연산 부분만 남기면 3배 정도는 빠르게 할 수 있다는 것이겠습니다 (뇌피셜).
from time import time
ort_time = time()
ort_output = ort_session.run(output_name, ort_inputs)
print("ONNX inference time:", time()-ort_time, "sec")
pt_time = time()
with torch.no_grad():
pt_output = cola_model(**input_sample)
print("PyTorch inference time:", time()-pt_time, "sec")
결과
ONNX inference time: 0.004585742950439453 sec
PyTorch inference time: 0.015786170959472656 sec
Netron
Netron은 모델의 구조를 보여주는 프로그램입니다. Tensorboard는 모델을 직접 구현했을 당시 그 구조를 보여주는 시각화 툴로 아주 유용하게 사용되어 왔습니다. 그런데 ONNX는 모델을 packaging 해버렸기 때문에 그 구조를 알아내기 어렵습니다. 코드를 직접 확인할 수가 없으니까요. 또한 input과 output의 이름들을 모르는 상태에서 onnx 파일만 가지고 있다면 사용하기 어렵겠죠. 이 때 Netron이라는 툴을 사용할 수 있습니다. .onnx, .pt 등 다양한 모델 파일들을 지원합니다. 단순 업로드만 해주면…!! 모델을 그려줍니다!! 각 layer의 이름 input, output의 형태 등 대부분의 정보를 다 알려준다고 보면 되겠습니다.
