Skip to main content

PyTorch Lightning과 Weights & Biases

PyTorch와 PyTorch Lightning을 Weights & Biases와 함께 비교하기 이 글은 AI 번역본입니다. 오역이 있을 경우 댓글로 알려 주세요.
Created on September 15|Last edited on September 15
PyTorch Lightning을 사용하면 연구용 코드와 엔지니어링 코드를 분리할 수 있습니다. 이 빠른 튜토리얼로 Lightning 모델을 시각화하고 Weights & Biases와의 손쉬운 통합을 통해 하이퍼파라미터를 최적화해 보세요.
PyTorch Lightning 사용해 보기 →, 또는 이 통합을 살펴보거나 실시간 대시보드 →.

설치 및 소개

PyTorch Lightning 설치는 매우 간단합니다:
pip install pytorch-lightning
우리의 PyTorch 코드에서 사용하려면, 필요한 PyTorch Lightning 모듈을 임포트합니다:
import pytorch_lightning as pl
from pytorch_lightning.loggers import WandbLogger
실험 결과를 추적하고 W&B에 직접 기록하기 위해 WandbLogger를 사용하겠습니다.

Lightning 클래스 만들기

연구를 하다 보면 실험적 변형을 추가하기 위해 보일러플레이트 코드를 자주 수정하게 됩니다. 이 과정에서 대부분의 오류가 코드베이스에 유입됩니다. PyTorch Lightning은 모델 정의와 학습을 위한 명확한 코드 구조를 제공하여 보일러플레이트 코드를 크게 줄여줍니다.
PyTorch에서 신경망 클래스를 만들려면 다음을 임포트하거나 상속해야 합니다 torch.nn.module마찬가지로 PyTorch Lightning을 사용할 때는 해당 클래스를 임포트합니다. pl.LightningModule.
이제 MNIST 데이터셋 분류 모델을 학습하기 위해 사용할 클래스를 만들어 보겠습니다. 결과를 비교하기 위해 공식 문서의 예제와 동일한 예시를 사용하겠습니다.
class LightningMNISTClassifier(pl.LightningModule):
def __init__(self):
super(LightningMNISTClassifier, self).__init__()

# mnist images are (1, 28, 28) (channels, width, height)

self.layer_1 = torch.nn.Linear(28 * 28, 128)
self.layer_2 = torch.nn.Linear(128, 256)
self.layer_3 = torch.nn.Linear(256, 10)
위에서 볼 수 있듯이, 임포트한 기본 클래스만 제외하면 코드의 나머지는 원래의 PyTorch 코드와 거의 동일합니다.
PyTorch에서는 이 데이터 로딩을 메인 학습 파일의 어느 위치에서나 수행할 수 있습니다.
PyTorch Lightning에서는 LightningModule의 세 가지 특정 메서드에서 이를 수행합니다.
  • train_dataloader()
  • val_dataloader()
  • test_dataloader()
그리고 데이터 준비 및 다운로드를 위한 네 번째 메서드가 있습니다.
  • prepare_data()
이렇게 하면 됩니다 코드에서.
def prepare_data(self):
# prepare transforms standard to MNIST
MNIST(os.getcwd(), train=True, download=True)
MNIST(os.getcwd(), train=False, download=True)

def train_dataloader(self):
#Load the dataset
mnist_train = DataLoader(self.mnist_train, batch_size=64)
return mnist_train

def val_dataloader(self):
#Load val dataset
mnist_val = DataLoader(self.mnist_val, batch_size=64)
return mnist_val

def test_dataloader(self):
#Load test data
mnist_test = DataLoader(mnist_test, batch_size=64)
return mnist_test
최적화기 코드는 Lightning에서도 동일하지만, 함수에 추가된다는 점만 다릅니다. configure_optimizers() LightningModule에서
전통적인 PyTorch 학습 파이프라인을 생각해 보면, 에포크 루프를 구현하고, 미니배치를 반복하며, 각 미니배치에 대해 순전파를 수행하고, 손실을 계산하고, 각 배치마다 역전파를 수행한 뒤 마지막으로 그래디언트를 업데이트해야 합니다.
Lightning에서 동일하게 하려면, 학습 루프와 검증 루프의 핵심 부분을 세 가지 함수로 분리합니다:
  • training_step
  • validation_step
  • validation_end
이 함수들의 프로토타입은 다음과 같습니다:
def training_step(self, train_batch, batch_idx):
def validation_step(self, val_batch, batch_idx):
def validation_end(self, outputs):
이 함수들을 사용하면 PyTorch Lightning이 파이프라인의 학습 과정을 자동화합니다. 그 내용은 곧 살펴보겠습니다. 그에 앞서, PyTorch Lightning이 Weights & Biases와 손쉽게 연동되어 실험을 추적하고 어디서든 모니터링할 수 있는 시각화를 만드는 방법부터 확인해 보겠습니다.

W&B로 PyTorch Lightning 모델 성능 추적하기

wandbLogger가 Lightning과 어떻게 연동되는지 살펴보겠습니다.
from pytorch_lightning.loggers import WandbLogger
wandb_logger = WandbLogger(name='Adam-32-0.001',project='pytorchlightning')
여기서는 프로젝트와 현재 기록 중인 런에 대한 정보를 담는 wandbLogger 객체를 생성했습니다.

훈련 루프

이제 어떤 모델이든 학습에서 가장 중요한 부분인 훈련 루프로 들어가 보겠습니다. 우리는 PyTorch Lightning을 사용하므로 대부분의 작업은 내부에서 자동으로 처리됩니다. 몇 가지 하이퍼파라미터만 지정하면 학습 과정이 자동으로 진행됩니다. 추가로, 각 반복마다 보기 좋은 진행 막대도 제공됩니다.
model = LightningMNISTClassifier()
model.prepare_data()
model.train_dataloader()
trainer = pl.Trainer(max_epochs = 5,logger= wandb_logger)
시각화와 관련해 코드에서 중요한 부분은 Lightning의 Trainer 객체에 logger로 WandbLogger 객체를 전달하는 부분입니다. 이렇게 하면 해당 로거가 자동으로 결과를 기록합니다.
def train():
trainer.fit(model)
Lightning으로 PyTorch 모델을 학습시키는 데 필요한 작업은 이것만으로 충분합니다. 이 한 줄의 코드가 복잡하고 비효율적인 기본 PyTorch 코드를 손쉽게 대체합니다.
PyTorch는 각 반복을 추적할 수 있는 보기 좋은 진행 막대도 제공합니다.


Weights & Biases로 성능 시각화하기

이번 실행에서 생성된 시각화를 살펴보겠습니다. 대시보드에서.

모델을 학습하는 동안 해당 실행의 학습 손실과 검증 손실이 대시보드에 실시간으로 자동 기록됩니다.
다른 실행을 비교하기 위해 동일한 학습 단계를 하이퍼파라미터만 바꿔 반복할 수 있습니다. 각 실행을 구별할 수 있도록 로거 이름을 변경하겠습니다.
wandb_logger = WandbLogger(name='Adam-32-0.001',project='pytorchlightning')
wandb_logger = WandbLogger(name='Adam-64-0.01',project='pytorchlightning')
wandb_logger = WandbLogger(name='sgd-64-0.01',project='pytorchlightning')
여기서는 실행 이름을 짓는 규칙을 사용했습니다. 첫 번째는 옵티마이저, 두 번째는 미니배치 크기, 세 번째는 학습률입니다. 예를 들어 ‘Adam-32-0.001’이라는 이름은 옵티마이저가 Adam이고 배치 크기가 32이며 학습률이 0.001임을 의미합니다.
지금까지 우리 모델의 성능은 다음과 같습니다.


이러한 시각화는 프로젝트에 영구적으로 저장되므로, 서로 다른 하이퍼파라미터로 만든 변형들의 성능을 비교하고, 최고 성능의 모델을 복원하며, 팀과 결과를 공유하는 일이 훨씬 쉬워집니다.

멀티 GPU 학습

Lightning은 데이터 병렬 처리와 멀티 GPU 학습을 위한 간단한 API를 제공합니다. 샘플러에서 torch의 데이터 병렬 클래스(예: DataParallel)를 직접 사용할 필요가 없습니다. 사용할 병렬 처리 모드와 GPU 개수만 지정하면 됩니다.
학습에는 여러 가지 방법이 있습니다.
  • 데이터 병렬(distributed_backend='dp') (다중 GPU, 1대의 머신)
  • DistributedDataParallel(distributed_backend='ddp') (여러 대의 머신에 걸친 다중 GPU).
  • DistributedDataParallel2(distributed_backend='ddp2') (한 머신 내부에서는 dp, 머신 간에는 ddp).
  • TPU(num_tpu_cores=8|x) (단일 TPU 또는 TPU 팟)
이 글에서는 데이터 병렬 백엔드를 사용하겠습니다. 기존 코드에 이를 통합하는 방법은 다음과 같습니다.
trainer = pl.Trainer(max_epochs = 5,logger= wandb_logger, gpus=1, distributed_backend='dp')
여기서는 Google Colab에서 작업하고 있으므로 GPU를 1개만 사용합니다.
GPU를 더 많이 사용할수록, 아래 그래프처럼 wandb에서 서로 다른 구성 간 메모리 사용량 차이를 모니터링할 수 있습니다.


조기 종료

PyTorch Lightning은 조기 종료를 적용하는 두 가지 방법을 제공합니다. 다음과 같이 사용할 수 있습니다:
A) early_stop_callback을 True로 설정하세요. 'val_loss'를 기준으로 확인합니다.

# in validation_end() return dict. If it is not found an error is raised.
trainer = Trainer(early_stop_callback=True)

B) 또는 직접 콜백을 구성하세요
early_stop_callback = EarlyStopping(
monitor='val_loss',
min_delta=0.00,
patience=3,
verbose=False,
mode='min')
trainer = Trainer(early_stop_callback=early_stop_callback)
우리가 만든 만큼 validation_end() function, 우리는 직접 설정할 수 있습니다 early_stop_callback = true:
trainer = pl.Trainer(max_epochs = 5,logger= wandb_logger, gpus=1, distributed_backend='dp',early_stop_callback=True)

16비트 정밀도

프로젝트의 요구 사항에 따라 모델 가중치의 정밀도를 높이거나 낮춰야 할 수 있습니다. 정밀도를 낮추면 더 큰 모델을 GPU 메모리에 맞출 수 있습니다. 이제 PyTorch Lightning에서 16비트 정밀도를 적용하는 방법을 알아봅시다.
먼저 NVIDIA Apex를 설치해야 합니다. 이를 위해 Colab에서 셸 스크립트를 만들어 실행하겠습니다.
%%writefile setup.sh
git clone https://github.com/NVIDIA/apex
pip install -v --no-cache-dir ./apex
!sh setup.sh
apex를 설치한 후에는 런타임을 다시 시작해야 합니다.
이제 Trainer의 precision 매개변수에 필요한 값을 직접 전달할 수 있습니다.
trainer = pl.Trainer(max_epochs=100, logger=wandb_logger, gpus=1, distributed_backend='dp', early_stop_callback=True, amp_level='O1', precision=16)

모델 저장 및 불러오기

연구를 하다 보면 모델을 일정한 구간으로 나누어 학습해야 할 때가 자주 있습니다. 이때 학습을 중지하고 상태를 저장한 뒤, 나중에 그 상태를 불러와 중단한 지점부터 학습을 재개해야 합니다.
모델을 저장하고 복원할 수 있으면 팀과 더 효율적으로 협업할 수 있고, 몇 주 전의 실험으로도 쉽게 돌아갈 수 있습니다.
To save pytorch lightning models with W&B, we use:
trainer.save_checkpoint('EarlyStoppingADam-32-0.001.pth')
wandb.save('EarlyStoppingADam-32-0.001.pth')
이렇게 하면 로컬 런타임에 체크포인트 파일이 생성되고, wandb로 업로드됩니다. 이제 다른 시스템에서 학습을 재개하기로 해도, wandb에서 체크포인트 파일을 불러와 프로그램에 다음과 같이 로드하면 됩니다:
wandb.restore('EarlyStoppingADam-32-0.001.pth')
model.load_from_checkpoint('EarlyStoppingADam-32-0.001.pth')
이제 체크포인트가 모델에 로드되었으므로, 원하는 학습 모듈을 사용해 학습을 재개할 수 있습니다.

PyTorch와의 비교

이제 PyTorch Lightning이 제공하는 간결한 프레임워크를 살펴보았으니, PyTorch와 어떻게 비교되는지 간단히 보겠습니다. Lightning에서는 Trainer를 생성한 뒤 train() 메서드를 호출하기만 하면 자동 콜백과 진행률 표시줄까지 포함해 모델을 학습할 수 있습니다.
이제 동일한 작업을 기본 PyTorch만으로 어떻게 구현하는지 살펴보겠습니다.
#Pytorch
pytorch_model = MNISTClassifier(
optimizer = torch.optim.Adam(pytorch_model.parameters(), lr=1e-3)

# ----------------
# LOSS
# ----------------

def cross_entropy_loss(logits, labels):
return F.nll_loss(logits, labels)

# ----------------
# TRAINING LOOP
# ----------------

num_epochs = 1
for epoch in range(num_epochs):

# TRAINING LOOP
for train_batch in mnist_train:
x, y = train_batch
logits = pytorch_model(x)
loss = cross_entropy_loss(logits, y)
print('train loss: ', loss.item())
loss.backward()

optimizer.step()
optimizer.zero_grad()

# VALIDATION LOOP

with torch.no_grad():
val_loss = []
for val_batch in mnist_val:
x, y = val_batch
logits = pytorch_model(x)
val_loss.append(cross_entropy_loss(logits, y).item())
val_loss = torch.mean(torch.tensor(val_loss))
print('val_loss: ', val_loss.item())

학습 코드가 얼마나 복잡해질 수 있는지 알 수 있습니다. 게다가 아직 멀티 GPU 학습, 조기 종료, Weights & Biases(wandb)로 성능을 추적하기 위한 수정 사항도 포함하지 않았습니다.
PyTorch에서 분산 학습을 추가하려면 데이터셋 샘플링을 위해 DistributedSampler를 사용해야 합니다.
def train_dataloader(self):
dataset = MNIST(...)
sampler = None

if self.on_tpu:
sampler = DistributedSampler(dataset)

return DataLoader(dataset, sampler=sampler)
조기 종료를 적용하려면 사용자 정의 함수를 직접 작성해야 합니다.
하지만 Lightning을 사용하면 이 모든 것을 코드 한 줄로 처리할 수 있습니다.
#Pytorch Lightning
trainer = pl.Trainer(max_epochs = 5,logger= wandb_logger, gpus=1, distributed_backend='dp',early_stop_callback=True)
trainer.fit(model)

이번 글은 여기까지입니다.

PyTorch Lightning 사용해 보기 →, 또는 이 통합을 살펴보거나 실시간 대시보드 →.


이 글은 AI로 번역된 기사입니다. 오역이 있을 수 있으니 댓글로 알려 주세요. 원문 보고서는 아래 링크에서 확인하실 수 있습니다: 원문 보고서 보기