torchscript

[yeomko22]

Reference

• https://pytorch.org/tutorials/beginner/Intro_to_TorchScript_tutorial.html
• https://pytorch.org/tutorials/advanced/cpp_export.html

example code

my_cell = MyCell()
x, h = torch.rand(3, 4), torch.rand(3, 4)
traced_cell = torch.jit.trace(my_cell, (x, h))

torchscript 사용 이유

• python interpreter가 아닌 자체 interpreter에 의해서 invoke 될 수 있음. 그래서 몇은 요청을 동시에 처리할 수 있음
• 모델을 통째로 디스크에 저장하고, 다른 환경에서 로드 할 수 있음. 예를들어 파이썬이 아닌 다른 언어로 짜여진 서버에서 로드 가능
• torchscript를 적용하면 코드 레벨에서 컴파일러 최적화 적용이 가능
• 많은 backend/device runtime 인터페이스를 제공한다.
• 즉, python이 배포에 적합한 언어는 아니다. 그러므로 python으로 개발된 모델을 python 의존성 없이 직렬화 하는 방법이 torchscript이다.

[yeomko22]

trace, script

• torchscript를 생성하는 방법으로는 trace와 script가 있다.
• trace의 경우 예시 인풋을 하나 넣어주고, 여기서 발생하는 연산을 추적해서 torchscript를 생성한다.
• 이 경우, 인풋 값에 따라 분기 문을 처리하는 control-flow를 torchscript에 반영하지 못한다.
• 만일 control-flow 반영이 필요할 경우에는 명시적으로 annotation을 달아서 torchscript compiler에게 전달해주는 script 기법을 사용해야 한다.

모델 직렬화, 파일로 저장

traced_script_module.save("traced_resnet_model.pt")

• 간단하게 save 메서드 호출만으로 torchscript를 직렬화 하여 파일로 저장이 가능하다.