Weights & Biases와 함께 Pytorch Lightning 사용하기

PyTorch Lightning은 PyTorch 코드를 구성하고 분산 훈련 및 16-비트 정밀도와 같은 고급 기능을 쉽게 추가할 수 있는 래퍼(lightweight wrapper)입니다

Weights & Biases 통합과 연동하여, 단 2줄의 코드 추가로 완전한 추적성 및 재현성을 위한 모델을 훈련 및 모니터링할 수 있습니다.

from pytorch_lightning.loggers import WandbLogger
wandb_logger = WandbLogger()

이 예시에서, 저희는 MNIST 데이터세트에 간단한 모델을 최적화하겠습니다.

Pytorch Lightning 해보기→

🚀 설치

Pytorch-lightning와 W&B는 pip를 통해 쉽게 설치할 수 있습니다.

pip install pytorch-lightning wandb

몇 개의Pytorch-lightning 모듈 및 WandbLogger만 불러오면 모델을 정의할 수 있습니다.

🏗️ LightningModule을 통해 모델 정의하기

연구는 흔히 새로운 실험 변형과 함께 보일러 플레이트 코드 편집을 포함합니다. 이러한 실험 프로세스(tinkering process)로 인해, 오류의 대부분은 코드베이스에 유입 됩니다. Pytorch lightning은 모델 정의 및 훈련을 위한 명확한 코드 구조(structurer)를 제공함으로써 보일러 플레이트 코드를 크게 감소시킵니다.

Pytorch에 신경망 네트워크 클래스(a neural network class)를 생성하려면, torch.nn.module에서 불러오거나 확장해야 합니다. 유사한 방식으로, Pytorch-Lightning을 사용할 경우, 저희는 클래스 pl.LightningModule를 사용합니다.

MNIST 데이터세트 분류를 위한 모델 훈련에 사용할 클래스를 생성해보겠습니다. 결과를 비교하기 위해 저희는 공식 설명서의 예시와 동일한 예시를 사용하겠습니다.

class LitMNIST(LightningModule):

  def __init__(self):
    super().__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에 상응할 것과 거의 동일합니다.

그런 다음, 몇 가지 추가적인 방법(methods) 정의할 필요가 있습니다.

• forward은 Pytorch에서처럼 추론에 사용됩니다
• training_step 은 forward와 유하사지만, 또한 단일 배치에서의 손실을 반환해야 하며, 이는 훈련 루프로써 자동으로 반복됩니다.

    def forward(self, x):
        '''method used for inference input -> output'''

        batch_size, channels, width, height = x.size()

        # (b, 1, 28, 28) -> (b, 1*28*28)
        x = x.view(batch_size, -1)
        x = self.layer_1(x)
        x = F.relu(x)
        x = self.layer_2(x)
        x = F.relu(x)
        x = self.layer_3(x)

        x = F.log_softmax(x, dim=1)
        return x

    def training_step(self, batch, batch_idx):
        '''needs to return a loss from a single batch'''
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)

        # Log training loss
        self.log('train_loss', loss)

        # Log metrics
        self.log('train_acc', self.accuracy(logits, y))

        return loss

옵티마이저의 정의는 Pytorch에서의 정의와 동일하나 configure_optimizes를 통해 수행되어야 합니다.

    def configure_optimizers(self):
        '''defines model optimizer'''
        return Adam(self.parameters(), lr=self.lr)

마지막으로 손실 및 메트릭 로깅을 위한validation_step및 test_step을 추가할 수 있습니다.

📊 데이터 로딩

다음과 함께 데이터 파이프라인을 생성할 수 있습니다:

• 🍦 Vanilla Pytorch DataLoaders
• ⚡ Pytorch Lightning DataModules

DataModules은 보다 더 체계적인 정의로, CPU 및 GPU간의 작업량 자동 분배와 같은 추가적인 최적화를 허용합니다.

DataModule은 또한 인터페이스에 의해 정의됩니다:

• prepare_data )은 GPU 하나에서 한 번만 호출됩니다. 일반적으로 아래의 데이터 다운로드 단계와 같습니다.
• setup은 각각의 GPU에서 개별적으로 호출되며 fit 또는 test단계인 경우, 정의할 스테이지를 받아들입니다.
• 각각의 데이터 세트를 로드할 train_dataloader, val_dataloader 및 test_dataloader

다음은 코드에 이 작업을 수행하는 방법입니다.

class MNISTDataModule(LightningDataModule):

    def __init__(self, data_dir='./', batch_size=256):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.transform = transforms.ToTensor()

    def prepare_data(self):
        '''called only once and on 1 GPU'''
        # download data
        MNIST(self.data_dir, train=True, download=True)
        MNIST(self.data_dir, train=False, download=True)

    def setup(self, stage=None):
        '''called one ecah GPU separately - stage defines if we are at fit or test step'''
        # we set up only relevant datasets when stage is specified (automatically set by Pytorch-Lightning)
        if stage == 'fit' or stage is None:
            mnist_train = MNIST(self.data_dir, train=True, transform=self.transform)
            self.mnist_train, self.mnist_val = random_split(mnist_train, [55000, 5000])
        if stage == 'test' or stage is None:
            self.mnist_test = MNIST(self.data_dir, train=False, transform=self.transform)

    def train_dataloader(self):
        '''returns training dataloader'''
        mnist_train = DataLoader(self.mnist_train, batch_size=self.batch_size)
        return mnist_train

    def val_dataloader(self):
        '''returns validation dataloader'''
        mnist_val = DataLoader(self.mnist_val, batch_size=self.batch_size)
        return mnist_val

    def test_dataloader(self):
        '''returns test dataloader'''
        mnist_test = DataLoader(self.mnist_test, batch_size=self.batch_size)
        return mnist_test

👟 모델 훈련

기존의 Pytorch 훈련 파이프라인을 고려하는 경우, 저희는 에포크에 대한 루프를 구현하고, 미니 배치를 반복하고, 각 미니 배치에 대한 피드 포워드 패스(feed forward pass)를 수행하고, 손실을 계산하며, 각 배치에 대한 역전파(backprop)을 수행하고 마지막으로 경사(gradient)를 업데이트해야 합니다.

Pytorch Lightning에서 동일한 작업을 수행하기 위해, 저희는 Pytorch Lightning 모듈 내의 훈련 로직과 데이터 로딩의 주요 요소를 추출했습니다.

이러한 함수를 사용하여, Pytorch Lighting은 파이프라인의 훈련 부분을 자동화합니다. Pytorch Lightning이 어떻게 Weights & Biases와 간단하게 통합되어 실험을 추적하고 어디서든 모니터링할 수 있는 시각화를 생성하는지에 대해서 우선 살펴보겠습니다.

WandB와 함께 Pytorch Lightning 모델 퍼포먼스 추적하기

wandbLogger가 어떻게 lightning과 통합되는지 살펴보겠습니다.

from pytorch_lightning.loggers import WandbLogger
wandb_logger = WandbLogger(name='Adam-32-0.001',project='pytorchlightning')

여기서, 저희는 프로젝트 및 로그 중인 실행에 대한 세부정보를 담은 wandbLogger 객체를 생성했습니다.

훈련 루프

Pytorching Lightning을 사용하고 있으므로, 대부분의 로직은 배후에 캡처되어 있습니다. 단지 몇몇 초매개변수만 지정하면 훈련 프로세스는 트레이너(Trainer)를 사용하여 자동으로 완료됩니다. 추가적인 이점으로, 각 반복에 대한 진행률 바도 제공됩니다.

전체 코드 보기 →

# setup data
mnist = MNISTDataModule()

# setup model - choose different hyperparameters per experiment
model = LitMNIST(n_layer_1=128, n_layer_2=256, lr=1e-3)

# define a Trainer
trainer = Trainer(
    logger=wandb_logger,    # W&B integration
    gpus=-1,                # use all GPU's
    max_epochs=3            # number of epochs
    )

시각화에 관한 코드에서 가장 중요한 부분은 WandbLogger 객체가 Pytorch Lightning의 트레이너(Trainer) 객체에서 로거(logger)로 전달된다는 것입니다. 이는 자동으로 로그를 사용하여 결과를 로그 합니다.

def train():
  trainer.fit(model)

이것이 Pytorch Lightning을 사용하여 Pytorch 모델을 훈련하기 위한 모든 것입니다. 이 한 줄의 코드는 크고 비효율적인 vanilla Pytorch 코드를 쉽게 대체합니다.

Weights & Biasesfmf 통한 퍼포먼스 시각화

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

특정 실행에 대한 훈련 손실 및 검증 손실은 모델을 훈련하는 동안 실시간으로 대시보드에 자동 로그 됩니다.

저희는 서로 다른 초매개변수와 함께 동일한 훈련 단계를 반복하여 서로 다른 실행을 비교할 수 있습니다. 저희는 각 훈련을 고유하게 식별하기 위해 로거의 이름을 변경하겠습니다.

전체 코드 보기 →

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')

여기서, 저는 실행의 이름을 지정하기 위해 컨벤션(convention)을 사용했습니다. 첫 번째 부분은 옵티마이저이며, 두 번째 부분은 미니배치 사이즈이고, 세 번째는 학습률입니다. 예를 들어, 이름‘Adam-32-0.001’는 사용 중인 옵티마이저는 배치 사이즈가 32인 Adam이고, 학습률은 0.001임을 의미합니다.

위의 플롯에서 각 모델이 퍼포먼스를 확인할 수 있습니다.

이러한 시각화는 프로젝트에 영구히 저장되며, 이는 변형(variations)의 퍼포먼스를 다른 초매개변수와 보다 쉽게 비교하고, 최고 성능의 모델을 복원하고, 팀원과의 결과 공유를 보다 쉽게 할 수 있도록 합니다.

Pytorch Lightning은 조기 종료(early stopping)을 포함하기 위해 2가지 방법을 제공합니다. 사용법은 다음과 같습니다.

전체 코드 보기 →

# A) Set early_stop_callback to True. Will look for 'val_loss'
# in validation_end() return dict. If it is not found an error is raised.
trainer = Trainer(early_stop_callback=True)

# B) Or configure your own callback
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)

검증 함수를 정의했으므로, early_stop_callback = true: 를 바로 설정할 수 있습니다.

trainer = pl.Trainer(max_epochs = 5,logger= wandb_logger, gpus=1, distributed_backend='dp',early_stop_callback=True)

프로젝트의 요구 사항에 따라, 모델의 가중치 정밀도를 높이거나 낮춰야 할 수 있습니다. 정밀도를 낮추면 더 큰 모델을 GPU에 맞출 수 있습니다. 16비트 정밀도를 pytorch lightning에 어떻게 포함하는지를 살펴보겠습니다.

먼저, NVIDIA apex를 설치해야 합니다. 이 작업을 수행하기 위해, 저희는 colab에 셸 스크립트(shell script)를 생성하여 실행합니다.

전체 코드 실행 →

%%writefile setup.sh

git clone https://github.com/NVIDIA/apex
pip install -v --no-cache-dir ./apex
!sh setup.sh 

apex 설치 후 런타임을 재시작 해야 합니다.

이제 트레이너의 정밀도 매개변수에 필수 값을 직접 전달할 수 있습니다.

trainer = pl.Trainer(max_epochs = 100,logger= wandb_logger, gpus=1, distributed_backend='dp',early_stop_callback=True, amp_level='O1',precision=16)

Lightning은 데이터 병렬 처리(Data parallelism) 및 다중 gpu 훈련을 수행하기 위한 간단한 API를 제공합니다. torch의 샘플러에서 데이터 병렬 처리 클래스를 사용할 필요가 없습니다. 병렬 모드 및 사용할 GPU의 숫자만 지정하시면 됩니다.

훈련 방법에는 다음의 여러 가지가 있습니다.

• Data Parallel (distributed_backend=’dp’) (다중 gpu, 머신 1대)

• DistributedDataParallel (distributed_backend=’ddp’) (여러 머신에서의 다중 gpu).

• DistributedDataParallel2 (distributed_backend=’ddp2’) (한 머신에는 dp, 여러 머신 간에는 ddp).

• TPUs (num_tpu_cores=8|x) (tpu 또는 TPU pod)

  이 포스팅에서는 데이터 병렬 백엔드를 사용합니다. 다음은 기존 코드에 통합하는 방법입니다.

trainer = pl.Trainer(gpus=1, distributed_backend='dp',max_epochs = 5,logger= wandb_logger)

Google colab에서 작업하는 동안 단 하나의 GPU만 사용합니다.

여러분은 더 많은 GPU를 사용하기 때문에, 왼쪽의 플롯에서처럼 wandb에서 다른 구성(configurations) 간 메모리 사용량 차이를 모니터링할 수 있습니다.

일반적으로 연구 중에는 간격을 두고 모델을 훈련하여야 합니다. 이는 훈련 중지, 상태 저장, 저장된 상태 로드 한 다음 중지한 지점에서 훈련을 재개해야 합니다.

모델을 저장 및 복원할 수 있으므로 팀원과 좀 더 효과적으로 협업할 수 있으며 몇 주 전의 실험으로 돌아갈 수 있습니다.

wandb를 통해 pytorch lightning을 저장하기 위해 저희는 다음을 사용합니다:

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')

이제 체크포인트가 모델로 로드 되었으며, 최적 훈련 모듈을 사용하여 훈련을 재개할 수 있습니다. Lightning이 제공하는 간소화한 프레임워크를 살펴봤음으로, 이제 Pytorch와 비교해보겠습니다.

Lighting에서, 트레이너 생성 및 train() 방법 호출하여 저희는 자동 콜백 및 진행률 바와 함께 모델을 훈련할 수 있습니다. Vanilla 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 훈련 통합, 조기 종료 또는 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 사용해보기 →..

Lightening과 Weights and Biases 통합에 관한 질문이 있으시면 저희 Slack 커뮤니티에 질문을 남겨주시기 바랍니다. 즐거운 마음으로 답변 드리겠습니다.