[PyTorch] TorchScript: Tracing vs. Scripting

안녕하세요, PyTorch로 학습한 모델을 배포하는 일은 굉장히 흔한 일인데요, 오늘은 TorchScript에서 활용하는 Tracing과 Scripting를 각각 설명하고 차이점을 비교하려고 합니다. 참고로 이 내용에는 여기 블로그가 많은 참고가 됐습니다.

그리고 TorchScript and PyTorch JIT | Deep Dive를 보시면 도움이 되실겁니다

시작에 앞서

두 방법은 모두 nn.Module을 전제로 합니다. PyTorch에서 nn.Module을 상속 받지 않은 모델은 존재하지 않는다고 봐도 무방합니다.

용어 설명

• Export: Eager-mode Python code를 그래프로 변환하는 과정
  (Eager-mode는 즉시 실행한다는 의미로, 컴파일을 하지 않는 파이썬의 기본 성질입니다)
• Tracing: Export의 한 방법으로, 임의의 Input을 모델에 넣었을 때 이루어지는 연산을 그대로 담는 방법
• Scripting: Python 코드를 컴파일하여 그래프를 얻는 방법
• TorchScript: 이건 두 가지 의미로 쓰일 수 있는데 - exported graph를 의미하거나, export 방법 그 자체를 의미합니다 (scripting) 여기서는 graph를 뜻하도록 하겠습니다
• Scriptable: torch.jit.script(model) 연산이 가능하면 scriptable이라고 함
• Traceable: torch.jit.trace(model, input) 연산이 가능하면 traceable이라고 함
• Dynamic control flow: input에 따라서 코드가 다르게 실행되는 경우를 의미. 예를 들어

# input x에 따라서 코드 분기점이 나뉘므로 Dynamic control flow
def func(x):
	if x[0] == 4:
		x += 1
    return x

# input x에 따라서 코드 분기점이 나뉘지 않는다! Dynamic control flow 아님
class A(nn.Module):
  backbone: nn.Module
  head: Optiona[nn.Module]
  def forward(self, x):
    x = self.backbone(x)
    if self.head is not None:
        x = self.head(x)
    return x

Scriptable 코드 작성의 의미

만약 누군가가 "우리는 파이썬 컴파일러를 만들어서 파이썬 개쩔게 만들겁니다"라고 말하면 그건 말도 안되는 소리에 가깝습니다. 왜냐면 파이썬은 그러기에는 너무 Dynamic하기 때문입니다.

 

그럼 혹시라도 파이썬 컴파일러를 만들 수 있다고 하면 그건 어떤 전제 조건이 붙을까요??

아주 기본적인(basic) syntax만으로 작성된 코드라면 그런게 가능할텐데, 그건 너무 어렵습니다...

대부분의 Python 코드는 방금 언급한 basic syntax 이외의 것들이 너무 많습니다. 예를 들어

custom structure, builtins, inheritance, Union, **kwargs, lambda, dynamic types

등이 존재합니다.

예를 들어 **kwargs를 쓰면 깔끔하게 쓸 수 있는 argument도 하나하나 나열해야 한다는 점에서 Scriptable은

코드를 다소 아름답지 못하게 만드는 것도 사실입니다.

Traceable 코드 작성의 의미

계속 작성하도록 하겠습니다!