Fork me on GitHub
Shortcuts

TorchScript 소개

Author: James Reed (jamesreed@fb.com), Michael Suo (suo@fb.com), rev2

번역: 강준혁

이 튜토리얼은 C++와 같은 고성능 환경에서 실행될 수 있는 PyTorch 모델(nn.Module 의 하위클래스)의 중간 표현인 TorchScript에 대한 소개입니다.

이 튜토리얼에서는 다음을 다룰 것입니다:

  1. 다음을 포함한 PyTorch의 모델 제작의 기본:

  • 모듈(Modules)

  • forward 함수 정의하기

  • 모듈을 계층 구조로 구성하기

  1. PyTorch 모듈을 고성능 배포 런타임인 TorchScript로 변환하는 특정 방법

  • 기존 모듈 추적하기

  • 스크립트를 사용하여 모듈을 직접 컴파일하기

  • 두 가지 접근 방법을 구성하는 방법

  • TorchScript 모듈 저장 및 불러오기

이 튜토리얼을 완료한 후에는 다음 학습서 를 통해 C++에서 TorchScript 모델을 실제로 호출하는 예제를 안내합니다.

import torch  # This is all you need to use both PyTorch and TorchScript!
print(torch.__version__)

Out:

1.2.0

PyTorch 모델 작성의 기초

간단한 Module 을 정의하는 것부터 시작하겠습니다. Module 은 PyTorch의 기본 구성 단위입니다. 이것은 다음을 포함합니다:

  1. 호출을 위해 모듈을 준비하는 생성자

  2. Parameters 집합과 하위 Module . 이것들은 생성자에 의해 초기화되며 호출 중에 모듈에 의해 사용될 수 있습니다.

  3. forward 함수. 모듈이 호출될 때 실행되는 코드입니다.

작은 예제로 시작해 보겟습니다:

class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()

    def forward(self, x, h):
        new_h = torch.tanh(x + h)
        return new_h, new_h

my_cell = MyCell()
x = torch.rand(3, 4)
h = torch.rand(3, 4)
print(my_cell(x, h))

Out:

(tensor([[0.9061, 0.6416, 0.4928, 0.6866],
        [0.9118, 0.9312, 0.6669, 0.8698],
        [0.9286, 0.6196, 0.8358, 0.8944]]), tensor([[0.9061, 0.6416, 0.4928, 0.6866],
        [0.9118, 0.9312, 0.6669, 0.8698],
        [0.9286, 0.6196, 0.8358, 0.8944]]))

우리는 다음 작업을 수행했습니다.:

  1. 하위 클래스로 torch.nn.Module 을 갖는 클래스를 생성했습니다.

  2. 생성자를 정의했습니다. 생성자는 많은 작업을 수행하지 않고 super 로 생성자를 호출합니다.

  3. 두 개의 입력을 받아 두 개의 출력을 반환하는 forward 함수를 정의했습니다. forward 함수의 실제 내용은 크게 중요하진 않지만, 가짜 RNN cell 의 일종입니다. 즉, 반복(loop)에 적용되는 함수입니다.

모듈을 인스턴스화하고, 3x4 크기의 무작위 값들로 이루어진 행렬 xy 를 만들었습니다. 그런 다음, my_cell(x, h) 를 이용해 cell을 호출했습니다. 이것은 forward 함수를 호출합니다.

좀 더 흥미로운 것을 해봅시다:

class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()
        self.linear = torch.nn.Linear(4, 4)

    def forward(self, x, h):
        new_h = torch.tanh(self.linear(x) + h)
        return new_h, new_h

my_cell = MyCell()
print(my_cell)
print(my_cell(x, h))

Out:

MyCell(
  (linear): Linear(in_features=4, out_features=4, bias=True)
)
(tensor([[ 0.5213, -0.4052,  0.4558,  0.2871],
        [ 0.5248, -0.0108,  0.6398,  0.8285],
        [ 0.6122,  0.0214,  0.7097,  0.7126]], grad_fn=<TanhBackward>), tensor([[ 0.5213, -0.4052,  0.4558,  0.2871],
        [ 0.5248, -0.0108,  0.6398,  0.8285],
        [ 0.6122,  0.0214,  0.7097,  0.7126]], grad_fn=<TanhBackward>))

모듈 MyCell 을 재정의했지만, 이번에는 self.linear 속성을 추가하고 forward 함수에서 self.linear 을 호출했습니다.

여기서 무슨 일이 일어날까요? torch.nn.LinearMyCell 과 마찬가지로 PyTorch 표준 라이브러리의 Module 입니다. 이것은 호출 구문을 사용하여 호출할 수 있습니다. 우리는 Module 의 계층을 구축하고 있습니다.

Module 에서 print 하는 것은 Module 의 하위 클래스 계층에 대한 시각적 표현을 제공할 것입니다. 이 예제에서는 Linear 의 하위 클래스와 하위 클래스의 매개 변수를 볼 수 있습니다.

Module 을 이런 방식으로 작성하면, 재사용 가능한 구성 요소를 사용하여 모델을 간결하고 읽기 쉽게 작성할 수 있습니다.

여러분은 출력된 내용에서 grad_fn 을 확인하셨을 것입니다. 이것은 오토그라드(autograd) 라 불리는 PyTorch의 자동 미분 방법의 세부 정보입니다. 요컨데, 이 시스템은 잠재적으로 복잡한 프로그램을 통해 미분을 계산할 수 있게 합니다. 이 디자인은 모델 제작에 엄청난 유연성을 제공합니다.

이제 유연성을 시험해 보겠습니다.

class MyDecisionGate(torch.nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x
        else:
            return -x

class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()
        self.dg = MyDecisionGate()
        self.linear = torch.nn.Linear(4, 4)

    def forward(self, x, h):
        new_h = torch.tanh(self.dg(self.linear(x)) + h)
        return new_h, new_h

my_cell = MyCell()
print(my_cell)
print(my_cell(x, h))

Out:

MyCell(
  (dg): MyDecisionGate()
  (linear): Linear(in_features=4, out_features=4, bias=True)
)
(tensor([[ 0.4222,  0.8903,  0.6804, -0.8304],
        [ 0.5213,  0.9497,  0.6797, -0.4223],
        [ 0.5283,  0.9206,  0.8842, -0.4677]], grad_fn=<TanhBackward>), tensor([[ 0.4222,  0.8903,  0.6804, -0.8304],
        [ 0.5213,  0.9497,  0.6797, -0.4223],
        [ 0.5283,  0.9206,  0.8842, -0.4677]], grad_fn=<TanhBackward>))

MyCell 클래스를 다시 정의했지만, 여기선 MyDecisionGate 를 정의했습니다. 이 모듈은 제어 흐름 을 활용합니다. 제어 흐름은 루프와 if 명령문과 같은 것으로 구성됩니다.

많은 프레임워크들은 주어진 프로그램 코드로부터 기호식 미분(symbolic derivatives)을 계산하는 접근법을 취하고 있습니다. 하지만, PyTorch에서는 변화도 테이프(gradient tape)를 사용합니다. 연산이 발생할 때 이를 기록하고, 미분값을 계산할 때 거꾸로 재생합니다. 이런 방식으로, 프레임워크는 언어의 모든 구문에 대한 미분값을 명시적으로 정의할 필요가 없습니다.

오토그라드가 작동하는 방식

오토그라드가 작동하는 방식

TorchScript의 기초

이제 실행 예제를 살펴보고 TorchScript를 적용하는 방법을 살펴보겠습니다.

한마디로, TorchScript는 PyTorch의 유연하고 동적인 특성을 고려하여 모델 정의를 캡쳐할 수 있는 도구를 제공합니다. 추적(tracing) 이라 부르는 것을 검사하는 것으로 시작하겠습니다.

Module 추적

class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()
        self.linear = torch.nn.Linear(4, 4)

    def forward(self, x, h):
        new_h = torch.tanh(self.linear(x) + h)
        return new_h, new_h

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

Out:

TracedModule[MyCell](
  (linear): TracedModule[Linear]()
)

살짝 앞으로 돌아가 MyCell 의 두 번째 버전을 가져왔습니다. 이전에 이것을 인스턴스화 했지만 이번엔 torch.jit.trace 를 호출하고, Module 을 전달했으며, 네트워크가 볼 수 있는 입력 예 를 전달했습니다.

여기서 무슨 일이 발생했습니까? Module 을 호출하였고, Module 이 돌아갈 때 발생한 연산을 기록하였고, torch.jit.ScriptModule 의 인스터스를 생성했습니다. ( TracedModule 은 인스턴스입니다)

TorchScript는 일반적으로 딥 러닝에서 그래프 라고 하는 중간 표현(또는 IR)에 정의를 기록합니다. .graph 속성으로 그래프를 확인해볼 수 있습니다:

print(traced_cell.graph)

Out:

graph(%self : ClassType<MyCell>,
      %input : Float(3, 4),
      %h : Float(3, 4)):
  %1 : ClassType<Linear> = prim::GetAttr[name="linear"](%self)
  %weight : Tensor = prim::GetAttr[name="weight"](%1)
  %bias : Tensor = prim::GetAttr[name="bias"](%1)
  %6 : Float(4!, 4!) = aten::t(%weight), scope: MyCell/Linear[linear] # /home/reserve/.pyenv/versions/3.6.7/envs/pytorch/lib/python3.6/site-packages/torch/nn/functional.py:1369:0
  %7 : int = prim::Constant[value=1](), scope: MyCell/Linear[linear] # /home/reserve/.pyenv/versions/3.6.7/envs/pytorch/lib/python3.6/site-packages/torch/nn/functional.py:1369:0
  %8 : int = prim::Constant[value=1](), scope: MyCell/Linear[linear] # /home/reserve/.pyenv/versions/3.6.7/envs/pytorch/lib/python3.6/site-packages/torch/nn/functional.py:1369:0
  %9 : Float(3, 4) = aten::addmm(%bias, %input, %6, %7, %8), scope: MyCell/Linear[linear] # /home/reserve/.pyenv/versions/3.6.7/envs/pytorch/lib/python3.6/site-packages/torch/nn/functional.py:1369:0
  %10 : int = prim::Constant[value=1](), scope: MyCell # /home/reserve/Workspace/PyTorch/TMPBuild/beginner_source/Intro_to_TorchScript_tutorial.py:184:0
  %11 : Float(3, 4) = aten::add(%9, %h, %10), scope: MyCell # /home/reserve/Workspace/PyTorch/TMPBuild/beginner_source/Intro_to_TorchScript_tutorial.py:184:0
  %12 : Float(3, 4) = aten::tanh(%11), scope: MyCell # /home/reserve/Workspace/PyTorch/TMPBuild/beginner_source/Intro_to_TorchScript_tutorial.py:184:0
  %13 : (Float(3, 4), Float(3, 4)) = prim::TupleConstruct(%12, %12)
  return (%13)

그러나, 이것은 저수준의 표현이며 그래프에 포함된 대부분의 정보는 최종 사용자에게 유용하지 않습니다. 대신, .code 속성을 사용하여 코드에 대한 Python 구문 해석을 제공할 수 있습니다:

print(traced_cell.code)

Out:

def forward(self,
    input: Tensor,
    h: Tensor) -> Tuple[Tensor, Tensor]:
  _0 = self.linear
  weight = _0.weight
  bias = _0.bias
  _1 = torch.addmm(bias, input, torch.t(weight), beta=1, alpha=1)
  _2 = torch.tanh(torch.add(_1, h, alpha=1))
  return (_2, _2)

어째서 이런 일들을 했을까요? 여기에는 몇 가지 이유가 있습니다:

  1. TorchScript 코드는 기본적으로 제한된 Python 인터프리터인 자체 인터프리터에서 호출될 수 있습니다. 이 인터프리터는 GIL(Global Interpreter Lock)을 얻지 않으므로 동일한 인스턴스에서 동시에 많은 요청을 처리할 수 있습니다.

  2. 이 형식을 사용하면 전체 모델을 디스크에 저장하고 Python 이외의 언어로 작성된 서버와 같은 다른 환경에서 불러올 수 있습니다.

  3. TorchScript는 보다 효율적인 실행을 제공하기 위해 코드에서 컴파일러 최적화를 수행할 수 있는 표현을 제공합니다.

  4. TorchScript를 사용하면 개별 연산자보다 프로그램의 더 넓은 관점을 요구하는 많은 백엔드/장치 런타임과 상호작용(interface)할 수 있습니다.

traced_cell 을 호출하면 Python 모듈과 동일한 결과가 생성됩니다:

print(my_cell(x, h))
print(traced_cell(x, h))

Out:

(tensor([[ 0.0947,  0.0470,  0.4898,  0.3312],
        [ 0.3073,  0.6995,  0.3745,  0.7623],
        [-0.1817,  0.5259,  0.5222,  0.4235]], grad_fn=<TanhBackward>), tensor([[ 0.0947,  0.0470,  0.4898,  0.3312],
        [ 0.3073,  0.6995,  0.3745,  0.7623],
        [-0.1817,  0.5259,  0.5222,  0.4235]], grad_fn=<TanhBackward>))
(tensor([[ 0.0947,  0.0470,  0.4898,  0.3312],
        [ 0.3073,  0.6995,  0.3745,  0.7623],
        [-0.1817,  0.5259,  0.5222,  0.4235]],
       grad_fn=<DifferentiableGraphBackward>), tensor([[ 0.0947,  0.0470,  0.4898,  0.3312],
        [ 0.3073,  0.6995,  0.3745,  0.7623],
        [-0.1817,  0.5259,  0.5222,  0.4235]],
       grad_fn=<DifferentiableGraphBackward>))

스크립팅을 사용하여 모듈 변환

제어 흐름이 포함된(control-flow-laden) 하위 모듈이 아닌 모듈 버전 2를 사용하는 이유가 있습니다. 지금 살펴봅시다:

class MyDecisionGate(torch.nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x
        else:
            return -x

class MyCell(torch.nn.Module):
    def __init__(self, dg):
        super(MyCell, self).__init__()
        self.dg = dg
        self.linear = torch.nn.Linear(4, 4)

    def forward(self, x, h):
        new_h = torch.tanh(self.dg(self.linear(x)) + h)
        return new_h, new_h

my_cell = MyCell(MyDecisionGate())
traced_cell = torch.jit.trace(my_cell, (x, h))
print(traced_cell.code)

Out:

def forward(self,
    input: Tensor,
    h: Tensor) -> Tuple[Tensor, Tensor]:
  _0 = self.linear
  weight = _0.weight
  bias = _0.bias
  x = torch.addmm(bias, input, torch.t(weight), beta=1, alpha=1)
  _1 = torch.tanh(torch.add(x, h, alpha=1))
  return (_1, _1)

.code 출력을 보면, if-else 분기가 어디에도 없다는 것을 알 수 있습니다! 어째서일까요? 추적은 코드를 실행하고 발생하는 작업을 기록하며 정확하게 수행하는 스크립트 모듈(ScriptModule)을 구성하는 일을 수행합니다. 불행하게도, 제어 흐름과 같은 것들은 지워집니다.

TorchScript에서 이 모듈을 어떻게 충실하게 나타낼 수 있을까요? Python 소스 코드를 직접 분석하여 TorchScript로 변환하는 스크립트 컴파일러(script compiler) 를 제공합니다. MyDecisionGate 를 스크립트 컴파일러를 사용하여 변환해 봅시다:

scripted_gate = torch.jit.script(MyDecisionGate())

my_cell = MyCell(scripted_gate)
traced_cell = torch.jit.script(my_cell)
print(traced_cell.code)

Out:

def forward(self,
    x: Tensor,
    h: Tensor) -> Tuple[Tensor, Tensor]:
  _0 = self.linear
  _1 = _0.weight
  _2 = _0.bias
  if torch.eq(torch.dim(x), 2):
    _3 = torch.__isnot__(_2, None)
  else:
    _3 = False
  if _3:
    bias = ops.prim.unchecked_unwrap_optional(_2)
    ret = torch.addmm(bias, x, torch.t(_1), beta=1, alpha=1)
  else:
    output = torch.matmul(x, torch.t(_1))
    if torch.__isnot__(_2, None):
      bias0 = ops.prim.unchecked_unwrap_optional(_2)
      output0 = torch.add_(output, bias0, alpha=1)
    else:
      output0 = output
    ret = output0
  _4 = torch.gt(torch.sum(ret, dtype=None), 0)
  if bool(_4):
    _5 = ret
  else:
    _5 = torch.neg(ret)
  new_h = torch.tanh(torch.add(_5, h, alpha=1))
  return (new_h, new_h)

만세! 이제 TorchScript에서 프로그램의 동작을 충실하게 캡쳐했습니다. 이제 프로그램을 실행해 봅시다:

# 새로운 입력
x, h = torch.rand(3, 4), torch.rand(3, 4)
traced_cell(x, h)

스크립팅과 추적 혼합

어떤 상황에서는 스크립팅보다는 추적을 사용해야 합니다. (예: 모듈에는 TorchScript에 표시하지 않으려는 Python 상수 값을 기반으로 만들어진 많은 구조적인 결정(architectural decisions)이 있습니다.) 이 경우, 스크립팅은 추적으로 구성될 수 있습니다: torch.jit.script 는 추적된 모듈의 코드를 인라인(inline) 할 것이고, 추적은 스크립트 된 모듈의 코드를 인라인 할 것입니다.

첫 번째 경우의 예:

class MyRNNLoop(torch.nn.Module):
    def __init__(self):
        super(MyRNNLoop, self).__init__()
        self.cell = torch.jit.trace(MyCell(scripted_gate), (x, h))

    def forward(self, xs):
        h, y = torch.zeros(3, 4), torch.zeros(3, 4)
        for i in range(xs.size(0)):
            y, h = self.cell(xs[i], h)
        return y, h

rnn_loop = torch.jit.script(MyRNNLoop())
print(rnn_loop.code)

Out:

def forward(self,
    xs: Tensor) -> Tuple[Tensor, Tensor]:
  h = torch.zeros([3, 4], dtype=None, layout=None, device=None, pin_memory=None)
  y = torch.zeros([3, 4], dtype=None, layout=None, device=None, pin_memory=None)
  y0, h0 = y, h
  for i in range(torch.size(xs, 0)):
    _0 = self.cell
    _1 = torch.select(xs, 0, i)
    _2 = _0.linear
    weight = _2.weight
    bias = _2.bias
    _3 = torch.addmm(bias, _1, torch.t(weight), beta=1, alpha=1)
    _4 = torch.gt(torch.sum(_3, dtype=None), 0)
    if bool(_4):
      _5 = _3
    else:
      _5 = torch.neg(_3)
    _6 = torch.tanh(torch.add(_5, h0, alpha=1))
    y0, h0 = _6, _6
  return (y0, h0)

두 번째 경우의 예:

class WrapRNN(torch.nn.Module):
    def __init__(self):
        super(WrapRNN, self).__init__()
        self.loop = torch.jit.script(MyRNNLoop())

    def forward(self, xs):
        y, h = self.loop(xs)
        return torch.relu(y)

traced = torch.jit.trace(WrapRNN(), (torch.rand(10, 3, 4)))
print(traced.code)

Out:

def forward(self,
    argument_1: Tensor) -> Tensor:
  _0 = self.loop
  h = torch.zeros([3, 4], dtype=None, layout=None, device=None, pin_memory=None)
  h0 = h
  for i in range(torch.size(argument_1, 0)):
    _1 = _0.cell
    _2 = torch.select(argument_1, 0, i)
    _3 = _1.linear
    weight = _3.weight
    bias = _3.bias
    _4 = torch.addmm(bias, _2, torch.t(weight), beta=1, alpha=1)
    _5 = torch.gt(torch.sum(_4, dtype=None), 0)
    if bool(_5):
      _6 = _4
    else:
      _6 = torch.neg(_4)
    h0 = torch.tanh(torch.add(_6, h0, alpha=1))
  return torch.relu(h0)

이러한 방식으로, 스크립팅과 추적은 상황에 따라서 따로 사용되거나, 함께 사용될 수 있습니다.

모델 저장 및 불러오기

TorchScript 모듈을 아카이브 형식으로 디스크에 저장하고 불러오는 API를 제공합니다. 이 형식은 코드, 매개 변수, 속성과 디버그 정보를 포함합니다. 이것은 그 아카이브가 완전히 별개의 프로세스로 로드할 수 있는 모델의 독립 표현임을 의미합니다. 랩핑 된 RNN 모듈을 저장하고 로드해 봅시다:

traced.save('wrapped_rnn.zip')

loaded = torch.jit.load('wrapped_rnn.zip')

print(loaded)
print(loaded.code)

Out:

ScriptModule(
  (loop): ScriptModule(
    (cell): ScriptModule(
      (dg): ScriptModule()
      (linear): ScriptModule()
    )
  )
)
def forward(self,
    argument_1: Tensor) -> Tensor:
  _0 = self.loop
  h = torch.zeros([3, 4], dtype=None, layout=None, device=None, pin_memory=None)
  h0 = h
  for i in range(torch.size(argument_1, 0)):
    _1 = _0.cell
    _2 = torch.select(argument_1, 0, i)
    _3 = _1.linear
    weight = _3.weight
    bias = _3.bias
    _4 = torch.addmm(bias, _2, torch.t(weight), beta=1, alpha=1)
    _5 = torch.gt(torch.sum(_4, dtype=None), 0)
    if bool(_5):
      _6 = _4
    else:
      _6 = torch.neg(_4)
    h0 = torch.tanh(torch.add(_6, h0, alpha=1))
  return torch.relu(h0)

보시다시피, 직렬화는 모듈 계층과 검사한 코드를 유지합니다. 또한 모델을 로드할 수 있습니다. 예를 들어, Python 없이 실행하기 위해 모델을 C++ 로 로드할 수 있습니다.

더 읽을거리

튜토리얼을 완료했습니다! 관련 데모를 보려면 TorchScript를 사용하여 기계 번역 모델을 변환하기 위한 NeurIPS 데모를 확인하십시오: https://colab.research.google.com/drive/1HiICg6jRkBnr5hvK2-VnMi88Vi9pUzEJ

Total running time of the script: ( 0 minutes 0.527 seconds)

Gallery generated by Sphinx-Gallery

Docs

Access comprehensive developer documentation for PyTorch

View Docs

Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

Resources

Find development resources and get your questions answered

View Resources