[PyTorch] PyTorch C++ API 이해하기

이전 글에서 PyTorch 내부의 Caffe2 라이브러리에서 어떻게 다양한 변수 타입들을 지원하고, 다양한 디바이스들을 지원하는지 확인했었습니다. 이번에 그보다 상위 레벨에서 PyTorch가 어떻게 동작하는 이해하기 위해서 C++ API에 대해서 알아보겠습니다.

PyTorch 내부 구조

위의 이미지를 보게 되면 많은 분들이 아시는거 처럼 PyTorch는 외부 인터페이스는 Python으로 되어 있고, 계산 기능에 해당하는 내부 구현(backend)은 대부분 C/C++로 되어 있습니다. (Python C++ extension)

 

Python은 C++에 비해 접근성이 편하다는 장점을 통해서 많은 분들이 PyTorch를 Python의 형태로 접하고 있습니다. 하지만 낮은 지연성이 요구되거나 멀티스레드 환경에서는 Python을 사용하는데 어려움이 있습니다. 따라서 이러한 환경에서는 C++ API를 사용해야 하기 때문에 지금 부터 차근차근 알아보도록 하겠습니다.

 

 

PyTorch C++ API는 크게 5가지로 구성되어 있습니다.

• ATen: 다른 모든 것의 기반이 되는 텐서 및 수학적 연산 라이브러리입니다.
• Autograd: 자동 미분 기능으로 ATen을 보강합니다.
• C++ Frontend: 머신러닝 모델의 훈련 및 평가를 위한 고수준 구조.
• TorchScript: TorchScript JIT 컴파일러 및 인터프리터에 대한 인터페이스.
• C++ Extensions: 사용자 지정 C++ 및 CUDA routines으로 Python API를 확장하는 수단.

(Class Hierarchy는 아래의 링크에서 확인 가능합니다.)

Library API — PyTorch main documentation

ATen

ATen은 PyTorch의 근본이 되는 라이브러리입니다. 따라서 Tensor Class (Tensors는 다차원 데이터)를 제공하고, 수백 가지의 Operation이 정의되어 있습니다. (aten에서는 Tensor에 대한 interface만 있고 실제 구현은 c10에 되어 있습니다) 그리고 이러한 대부분의 Operation들은 CPU와 GPU가 모두 구현되어 있으며, 각각의 Operation에서는 Tensor 클래스의 유형에 따라 다르게 동작하게 됩니다. 그래서 만약 어떤 Kernal에 대한 찾고 있다면 ATen쪽 코드베이스에서 먼저 확인해 보는것이 좋습니다.

(참고로, ATen은 "A Tensor Library"의 줄임말 입니다.)

Tensor의 구조

 

 

추가적으로 현재 ATen의 core는 천천히 c10으로 이동 및 리펙토링중입니다. 기존에 aten이 존재했고 이를 사용하다가 caffe2가 PyTorch에 병합되면서 기본 동작들을 c10이라는 새로운 폴더에 이동 및 리펙토링을 진행하는 것으로 보입니다. 따라서 추후 마이그레이션이 마무리 되면 ATen에는 Operation들만 남고, Tensor와 Storage에 대한 구현과 PyTorch의 기본동작은 c10에 모아지게 될 것으로 보입니다.

Autograd

Autograd는 것은 자동 미분과 관련된 기능으로 ATen의 Tensor 클래스를 강화하는 C++ API 입니다.
autograd 시스템은 텐서에 대한 작업을 기록하여 autograd 그래프를 형성합니다. 그리고, backwards() 가 호출되면 기존에 기록한 그래프를 통해 backpropagation을 수행하게 되고 궁극적으로 미분 값을 얻게 됩니다.

사용할 일은 거의 없지만 아래의 같은 간단한 예제로 미분의 동작을 확인할 수 있습니다.

#include <torch/csrc/autograd/variable.h>
#include <torch/csrc/autograd/function.h>

torch::Tensor a = torch::ones({2, 2}, torch::requires_grad());
torch::Tensor b = torch::randn({2, 2});
auto c = a + b;
c.backward(); // a.grad() will now hold the gradient of c w.r.t. a.

ATen의 at::Tensor 클래스는 기본적으로 미분 가능하지 않습니다. 따라서 autograd API가 제공하는 Tensor의 미분을 가능하게 하려면 at:: 네임스페이스 대신 torch:: 네임스페이스의 Tensor Factory 함수를 사용해야 합니다.   예를 들어, at::ones로 생성된 Tensor는 미분 가능하지 않지만, torch::ones로 생성된 Tensor 는 미분 가능합니다.

 

C++ Frontend

C++ Frontend는 기계 학습 및 신경망에 대한 자동 미분 및 다양한 상위 수준 추상화 제공하며, 주로 Python API 설계 및 제공 기능을 따릅니다. 따라서 일부 환경에서 높은 성능(게임, 서버, 멀티 스레드 환경 등등..) 혹은 이식성 요구 사항으로 인해 Python 인터프리터를 사용할 수 없는 경우 (Python의 Frontend 사용자 경험을 희생하지 않으면서) 사용할 수 있습니다. 

따라서 C++ Frontend는 Python의 높은 사용성을 따라하기 위해 경우에 따라서 성능을 희생했다고 합니다. 

PyTorch에서 Python이 C++ 보다 항상 느린것은 아닙니다. Python Frontend의 계산이 많이드는 거의 모든 작업을 C++로 호출하고 있기 때문에, Python 작업을 선호하고 위에서 언급한 제약된 환경이 아니라면 Python을 사용하여 작업하라고 이야기하고 있습니다.

 

C++ Frontend 구체적으로 다음 구성 요소로 구성됩니다.

또한 아래의 기능들을 포함하고 있습니다.

• 계층적 모듈 시스템(예: torch.nn.Module)을 통해 기계 학습 모델을 정의하기 위한 인터페이스입니다.
• 가장 일반적인 모델링 목적(convolutions, RNN, batch normalization 등)을 위한 기존 모듈의 "표준 라이브러리"입니다.
• SGD, Adam, RMSprop 등과 같은 널리 사용되는 optimizers의 구현을 포함한 최적화 API
• 많은 CPU 코어에 걸쳐 데이터를 병렬로 로드하는 기능을 포함하여 데이터 세트 및 데이터 파이프라인을 표현하는 수단입니다.
• 훈련 세션(torch.utils.data.DataLoader)의 체크포인트를 저장하고 로드하기 위한 직렬화 형식
• 모델을 여러 GPU로 자동 병렬화(torch.nn.parallel.DataParallel)
• pybind11을 사용하여 C++ 모델을 Python에 쉽게 바인딩하는 코드를 지원
• TorchScript JIT 컴파일러에 대한 진입점
• ATen 및 Autograd API와의 인터페이스를 용이하게 하는 유용한 유틸리티입니다.

위 기능들을 보면 알 수 있듯이 Python을 통해 PyTorch을 사용했었던 것과 비슷한 인터페이스로 C++에서 사용할 수 있습니다. 아래의 예제는 MNIST 예제를 분류하는 모듈을 정의하고 학습하고 저장하는 ene-to-end 기능들을 C++로 구현한 예제입니다.

#include <torch/torch.h>

// 모듈 정의
struct Net : torch::nn::Module {
  Net() {
    // 3개의 Linear 하위 모듈을 구성하고 등록
    fc1 = register_module("fc1", torch::nn::Linear(784, 64));
    fc2 = register_module("fc2", torch::nn::Linear(64, 32));
    fc3 = register_module("fc3", torch::nn::Linear(32, 10));
  }

  // Net 알고리즘 구현
  torch::Tensor forward(torch::Tensor x) {
    // 다양한 Tensor 조작 함수 중 하나를 사용
    x = torch::relu(fc1->forward(x.reshape({x.size(0), 784})));
    x = torch::dropout(x, /*p=*/0.5, /*train=*/is_training());
    x = torch::relu(fc2->forward(x));
    x = torch::log_softmax(fc3->forward(x), /*dim=*/1);
    return x;
  }

  // 많은 "표준 라이브러리" 모듈 중 하나를 사용
  torch::nn::Linear fc1{nullptr}, fc2{nullptr}, fc3{nullptr};
};

int main() {
  // Net 생성
  auto net = std::make_shared<Net>();

  // MNIST dataset을 위한 멀티스레드 데이터로더 생성
  auto data_loader = torch::data::make_data_loader(
      torch::data::datasets::MNIST("./data").map(
          torch::data::transforms::Stack<>()),
      /*batch_size=*/64);

  // Net의 파리미터들을 업데이트하기 위해서 SGD optimization 알고리즘을 생성
  torch::optim::SGD optimizer(net->parameters(), /*lr=*/0.01);

  for (size_t epoch = 1; epoch <= 10; ++epoch) {
    size_t batch_index = 0;
    // 데이터 로더를 반복하여 데이터 세트에서 배치를 생성
    for (auto& batch : *data_loader) {
      // 미분 초기화
      optimizer.zero_grad();
      // 입력 데이터에 대해 모델을 실행
      torch::Tensor prediction = net->forward(batch.data);
      // 모델의 예측을 판단하기 위해 손실 값을 계산
      torch::Tensor loss = torch::nll_loss(prediction, batch.target);
      // 우리 모델의 매개변수 w,r, t에 대해서 loss의 기울기를 계산
      loss.backward();
      // 계산된 기울기를 기반으로 매개변수를 업데이트
      optimizer.step();
      // 100개의 배치마다 loss와 checkpoint를 출력
      if (++batch_index % 100 == 0) {
        std::cout << "Epoch: " << epoch << " | Batch: " << batch_index
                  << " | Loss: " << loss.item<float>() << std::endl;
        // 모델을 주기적으로 체크포인트로 직렬화
        torch::save(net, "net.pt");
      }
    }
  }
}

 

이러한 예제를 구현하기 위해서 직접 PyTorch를 빌드해볼 수도 있겠지만, 특별한 경우(자신만의 특별안 Op를 적용하거나 특별한 하드웨어 동작을 추가한 경우)가 아니라며 이 링크에서 배포판을 받아서 간단하게 적용해 볼 수 있습니다. 링크에서 CMake 구성과 간단예제가 포함되어 있어서 적용하는데 크게 어려움은 없습니다.

 

이번에는 PyTorch의 기본적인 동작을 구성하는 C++ API에 대해서 살펴보았습니다. 다음 글에서는 Custom Op를 추가하면서, PyTorch의 확장성에 해당하는 TorchScript와 C++ Extension에 대해서 이야기 해보겠습니다.

 

 

REFERENCE

PyTorch C++ API — PyTorch main documentation

PyTorch C++ API — PyTorch main documentation

PyTorch (tudelft.nl)

PyTorch

PyTorch internals : ezyang’s blog

PyTorch internals : ezyang’s blog