본문 바로가기
인공지능

Tensor에 대하여(with PyTorch, TPU)

hot6 2024. 12. 8. 21:57

개요

이 글은 PyTorch를 기반으로 DQN 관련 프로젝트를 하면서, 이 프레임워크에 대해 가볍게 정리해보고 싶었던 것들에 대한 글이다

 

우선 그동안 막연하게만 알았던 Tensor라는 데이터 타입에 대해 정리해본다

그리고 PyTorch에서 이런 Tensor들로 구성된 모델을 저장하고 불러오는 방법에 대해 정리한다

이런 모델은 CPU, GPU 뿐만 아니라 TPU에도 로드가 가능한데, TPU에 대해서도 간단히 알아본다

 

  • Tensor
  • 모델 저장 및 불러오기
  • TPU

 

Tensor

PyTorch의 핵심 개념 중 하나는 Tensor다

 

Tensor는 수학적으로는 다차원 배열을 의미하며, PyTorch에서 데이터를 저장하고 연산을 수행하는 기본 자료 구조다

NumPy의 배열과 상당히 유사하여 상호변환 시, 오버헤드가 거의 없다(메모리를 공유하기 때문)

 

아래와 같이 일반적인 Python의 배열은 데이터의 주소를 가리키는 요소들로 이뤄진다

즉, C의 배열과 다르게 데이터들이 순차적으로 메모리상에 위치하지 않는다

반면 NumPy 배열은 마치 C의 배열처럼 데이터들이 메모리상에 위치한다

 

 

https://jakevdp.github.io/PythonDataScienceHandbook/02.01-understanding-data-types.html

 

 

NumPy 배열을 자세히보면 dimensions와 strides 같은 요소가 있는데, 이는 메모리에 저장된 데이터들에 접근하는 방식에 사용된다

Tensor metadata

Tensor는 다음과 같이 Storage와 View라는 메타데이터를 가진다

Storage에 기록되어 있는 데이터를 View를 통해 볼 수 있다

Storage

  • 텐서가 실제로 데이터를 저장하는 공간
  • 텐서가 특정 차원에서 크기를 변경하거나 뷰를 변경해도 스토리지는 변하지 않는다

View

  • 텐서는 다양한 뷰를 생성 가능
  • 뷰는 동일한 데이터를 참조하면서도 다른 형태(reshape)로 표현

→ (3, 4) 크기의 텐서를 (12,) 또는 (4, 3)으로 표현(View)해도 실제 데이터는 동일한 메모리(Stage)를 참조한다

예시) Storage와 View의 이해

  • Tensor의 stride는 메모리에서 데이터를 "한 차원에서 다음 차원으로 이동"하는 단계(step)를 나타냄

 

파이토치 딥러닝 마스터(엘리 스티븐스, 루카 안티가, 토마스 피이만) ❘ 책만

 

위 사진처럼 전치 메서드를 사용하면, Storage 데이터의 순서가 바뀌는게 아니라 stride의 변경을 통해 데이터를 표현하는 방식이 달라질 뿐이다

 

즉, View는 마구잡이로 생성하더라도 저장되는 데이터의 순서에 변화는 없지만, 만약 그 데이터를 바꿔버리면 같은 Storage의 데이터가 바뀌기에 주의가 필요하다

Tensor operator

수학적 연산이 벡터화(vectorization) 되어 있어 반복문 없이 고속 수행 가능하다(행렬 곱셈, 미분, 통계 연산 등)

또한 broadcasting을 통해 서로 크기가 다른 텐서 연산 시 자동으로 크기를 맞춘다

다음과 같이 squeeze, unsqueeze 메서드는 딥러닝에서 배치(batch) 또는 채널(channel)을 조작할 때 유용하다

squeeze

크기가 1인 차원을 제거

예: (1, 3, 1, 5) → squeeze() → (3, 5)

unsqueeze

지정된 위치에 크기 1인 새로운 차원을 추가

예: (3, 5) → unsqueeze(0) → (1, 3, 5)

Device 설정

PyTorch에서 device를 통해 텐서가 저장되거나 연산이 수행되는 하드웨어를 지정 가능하다

CPU와 GPU 뿐만 아니라 TPU도 사용가능한데, 이때는 XLA(Accelerated Linear Algebra) 백엔드를 사용한다

 

아래처럼 Tensor 생성 시 device를 지정하거나, to 메서드를 통해 이동시킬 수 있다

 

import torch

# 사용 가능 디바이스 확인
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 텐서를 디바이스에 할당
tensor = torch.tensor([1, 2, 3], device=device)

# 기존 텐서를 디바이스로 이동
tensor_cpu = torch.tensor([4, 5, 6])
tensor_gpu = tensor_cpu.to(device)

 

모델 저장 및 불러오기

학습 모델을 다음과 같이 정의했다면, 모델의 파라미터만 저장하는 방식과 전체 모델의 구조를 저장하는 방식으로 모델을 저장할 수 있다

 

import torch

# 간단한 모델 정의
class SimpleModel(torch.nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 1)

    def forward(self, x):
        return self.fc(x)

State Dictionary 저장 및 불러오기

모델 구조는 저장되지 않고, 가중치(학습된 값)만 저장하는 방식

  • .state_dict()
  • .load_state_dict()

 

model = SimpleModel()

# 모델의 state_dict 저장
torch.save(model.state_dict(), "model_state.pth")

##### ...

# SimpleModel과 구조는 같으나 클래스명(SimpleModel2)은 달라도 불러오기 가능
loaded_model = SimpleModel2()

# state_dict 로드
loaded_model.load_state_dict(torch.load("model_state.pth"))

전체 모델 저장 및 불러오기

모델의 구조(아키텍처)와 파라미터를 함께 저장하는 방식

저장 당시 사용한 모델 클래스 정의를 그대로 유지해야 한다!

  • .save()
  • .load()

 

model = SimpleModel()

# 모델 저장
torch.save(model, "model.pth")

##### ...

# 모델 로드
loaded_model = torch.load("model.pth")

체크포인트 저장 및 불러오기

중단된 학습을 재개하기 위해 모델의 가중치뿐 아니라 Optimizer의 상태 및 학습 상태를 저장하고 불러와야 한다

 

  • Optimizer 상태
    • 모든 학습된 파라미터(모델의 가중치 및 편향).
    • 모멘텀, 학습률 스케줄러 등 추가 메타데이터.
  • 학습 상태
    • Epoch 및 Loss 상태

 

모델의 가중치와 학습상태들을 함께 저장하려면, dictionary 형태로 저장하는 것이 일반적이다

 

model = SimpleModel()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# Checkpoint 저장
checkpoint = {
    "model_state_dict": model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
    "epoch": epoch,  # 현재 epoch 정보도 저장
    "loss": loss  # 현재 loss 값도 저장
}

torch.save(checkpoint, "checkpoint.pth")

##### ...

# Checkpoint 로드
checkpoint = torch.load("checkpoint.pth")

# 모델과 Optimizer 상태 복원
model.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])

# 추가적인 학습 상태 복원
start_epoch = checkpoint["epoch"]
loss = checkpoint["loss"]

모델 추론

모델의 .eval() 메서드를 이용하면 추론 모드로 모델을 설정한다

이를 통해 드롭아웃이나 배치 정규화같은 레이어가 학습모드로 동작 안시킬 수 있다

또한 with torch.no_grad(): 컨텍스트를 통해 계산 그래프 생성을 비활성시킬 수 있다

 

model = torch.load("model.pth")
model.eval()

# input = ...

with torch.no_grad():
    outputs = model(inputs)

 

TPU(Tensor Processing Unit)

TPU(Tensor Processing Unit)는 Google이 설계된 하드웨어로, 특히 텐서 연산(Matrix Multiplication)에 특화되어 있다

PyTorch는 TPU를 지원하기 위해 PyTorch/XLA라는 확장 라이브러리를 사용하며, 이를 통해 PyTorch 코드를 변경하지 않고도 TPU에서 실행할 수 있도록 한다

XLA(Accelerated Linear Algebra)

Google이 개발한 컴파일러로, 딥러닝 모델의 연산을 최적화하여 하드웨어(GPU, TPU)에서 빠르고 효율적으로 실행할 수 있도록 설계되어 있다

 

  • 연산 그래프 최적화: 여러 연산을 하나로 합치거나, 불필요한 계산을 제거.
  • 메모리 최적화: 메모리 대역폭을 효율적으로 사용.
  • 배치 처리: TPU의 대규모 연산 코어를 최대한 활용.

PyTorch/XLA

PyTorch의 기존 API를 확장하여, PyTorch 코드 변경 없이 TPU를 활용할 수 있도록 한다

TPU 특화 연산인 Systolic Array 구조를 활용해 효율적인 계산 수행한다

이에 관해서는 간단하게 다음 블로그에서 gif를 보면 이해할 수 있는데, 나중에 NPU에 대해 한번 정리해 봐야겠다

 

마무리

이번 글에서는 PyTorch로 코드를 작성하며 기본 데이터 타입으로 사용되는 Tensor에 대해 중점적으로 살펴보았다

 

Tensor의 device로 GPU 말고 TPU도 설정할 수 있는데, TPU에서 모델을 실행하려면 XLA를 통해 모델 컴파일이 필요하다

TPU의 동작 방식을 살펴보다가, Systolic Array라는 설계 구조를 알게 되었는데, 이는 연산 코어들이 데이터를 선형적으로 전달하며 병렬 처리하는 방식이다

TPU를 포함한 NPU들은 이런 방식으로 데이터를 처리할 수 있는데, 애플 실리콘의 뉴럴 엔진도 이런 방식이겠구나라는 생각이 들었다

티스토리툴바