Pytorch-Lightning Commom Use Cases 05 - Inference in Production

한건우·2021년 10월 21일
0
  • Pytorch-Lightning은 production level의 model deploying을 간단하게 사용 가능함

ONNX로 export

  • Pytorch-Lightning은 model을 ONNX format으로 빠르게 export하는 간편한 기능을 제공함
filepath = "model.onnx"
model = SimpleModel()
input_sample = torch.randn((1, 64))
model.to_onnx(filepath, input_sample, export_params=True)
  • example_input_array를 skip 가능함
  • model을 export 한 경우 ONNX runtime에서 다음과 같이 사용가능
ort_session = onnxruntime.InferenceSession(filepath)
input_name = ort_session.get_inputs()[0].name
ort_inputs = {input_name: np.random.randn(1, 64).astype(np.float32)}
ort_outs = ort_session.run(None, ort_inputs)

TorchScript로 export

  • TorchScript는 non-Python 환경에서 불러 사용할 수 있도록 Selialize된 모델임
  • LightningModule은 to_torchscript() 로 바로 사용할 수 있도록 method를 제공함
model = SimpleModel()
script = model.to_torchscript()
torch.jit.save(script, "model.pt")
profile
아마추어 GAN잽이

0개의 댓글