PyTorch와 비교 해본 PyTorch Lightning 소개 (An introduction to PyTorch Lightning with comparisons to PyTorch)

이 포스트는 amaarora의 포스팅을 번역(및 내용추가)한 글입니다.

 

PytorchLightning을 이미 사용해보셨나요? 그렇다면, 당신은 그게 왜 쿨한지 아실거에요. 사용해보지 않으셨다면, 이 포스팅을 모두 일은 후 그것이 얼마나 쿨한지 아셨으면 좋겠어요(여기서 말하는 '그것'은 이 블로그 포스트나 멋진 PytorchLightning를 말하는 것입니다. 결정은 독자분들께 맡기겠습니다).

 

Note: 여기서부터 우리는 PytorchLightning를 PL이라고 줄여부를거에요, 타이핑하기엔 너무 길고 내가 좋아하는 키보드를 회사에 두고 왔기 때문입니다.

 

잠시동안, 저는 CPU, GPU, TPU로 모델을 많이 바꾸지않고 동일한 스크립트로 학습시킬 수 있기 때문에 Tensorflow를 질투했어요! 예를 들어, 이 제가 좋아하는 캐글러인 NVIDIA의 리서처이고 Kaggle 4x Grandmaster인 Chris Deotte의 notebook을 가져와서 블로그 포스트를 작성할 때요. Tensorflow에서 단지 적절한 strategy만 사용한다면, 같은 실험을 동일한 스크립트로 TPU, GPU, CPU에서 모두 실행해 볼 수 있습니다.

 

당신이 이미 멀티-GPU를 사용하거나 PyTorch로 TPU를 사용하기 위해 torch XLA를 사용한다면, 제가 말하고 있는게 뭔지 아실거에요. PyTorch에서 하드웨어(CPU, GPU, TPU)를 고르는 것은 불편합니다. 나는 PyTorch를 사랑해요 - 하지만 이것 하나 만큼은 정말 실망스러워요.

 

PL에 오신 것을 환영합니다! 이 라이브러리를 더 사용해보길 바라요.

 

이 포스트에서, 우리는 PL에 대한 소개와 몇 가지 쿨한 트릭을 사용하는 방법을 이야기할겁니다 - Gradient Accumulation,16-bit precision training,TPU/multi-gpu support - 같은 것 들을 몇 줄의 코드라인과 함께요. 우리는 PL을 캐글의 SIIM-ISIC Melanoma Classification를 해보면서 사용할거에요. 이 포스트에선 우리의 포커스는 PL의 소개와 ISIC 컴피티션을 예로 사용하는 것입니다.

 

우리는 그리고 PyTorch의 전형적인 워크플로우와 PL이 어떻게 다른지 알아볼거고, 이게 리서처들에게 얼마나 도움이 되는지 알아볼거에요.

 

블로그 포스트의 첫 번째는, data를 얻는 법일 것입니다. Train과 validation 데이터셋과 dataloader를 만들고 The Lightning Module섹션에서 PL의 흥미로운 부분들을 다룰거에요. 이 부분을 이미 여러번 해봐서 재미없는 분들은 바로 모델 구현으로 건너 뛰셔도 좋습니다.

 

1. What’s ISIC Melanoma Classification challenge?
2. Getting the data
3. Melanoma Dataset
4. Lightning Module
   1. Model and Training
   2. Model implementation compared to PyTorch
5. Gradient Accumulation
6. 16-bit precision training
7. TPU Support
8. Conclusion
9. Credits

ISIC Melanoma Classification challenge가 뭘까?

캐글의 챌린지 description를 보자면,

Skin cancer is the most prevalent type of cancer. Melanoma, specifically, is responsible for 75% of skin cancer deaths, despite being the least common skin cancer. The American Cancer Society estimates over 100,000 new melanoma cases will be diagnosed in 2020. It’s also expected that almost 7,000 people will die from the disease. As with other cancers, early and accurate detection—potentially aided by data science—can make treatment more effective. Currently, dermatologists evaluate every one of a patient’s moles to identify outlier lesions or “ugly ducklings” that are most likely to be melanoma.

이 컴피티션에서, 참가자들은 병변과 흑색종(Melanoma)을 분류하는 분류기를 만들어야합니다. 전형적인 병변은 아래의 이미지들과 같습니다.

src: https://www.isic-archive.com/#!/topWithHeader/onlyHeaderTop/gallery

이 포스트에서, 우리는 흑색종 분류 모델을 PL로 구현할 것입니다. 모델은 train 몇 시간만에 0.92 AUC 스코어를 낼 수 있습니다!

A side note: 딥러닝은 먼 길을 걸어 왔습니다. 2012년 AlexNet이 겨우 3GB 메모리를 가진 여러개의 GTX 580 GPU로 학습된 것가 비교한다면요. 120만 예를 가진 Imagenet은, 논문의 저자가 모델(레이어가 8개 뿐인)을 2개의 GPU에 나눠야만 했습니다. 그리고 Train하는 데 5-6일이 걸렸습니다. 오늘날, 우리는 모델을 train하는데 몇 시간, 몇 분이면 가능합니다. ISIC는 validation을 포함한 256X256 사이즈의 이미지를 각 에폭마다 P100 GPU로 2분이면 train 할 수 있습니다.

Data 가져오기

당신은 256x256 버전의 Jpeg 데이터를 여기서 가져올 수 있습니다.

Melanoma Dataset

모델 train에 사용할 데이터를 수집 할 준비를하는 것은 모든 프로젝트에서 수행해야하는 기본 작업 중 하나입니다.

class MelonamaDataset:
    def __init__(self, image_paths, targets, augmentations=None):
        self.image_paths = image_paths
        self.targets = targets
        self.augmentations = augmentations

    def __len__(self): return len(self.image_paths)

    def __getitem__(self, idx):
        image_path = self.image_paths[idx]
        image = np.array(Image.open(image_path))
        target = self.targets[idx]

        if self.augmentations is not None:
            augmented = self.augmentations(image=image)
            image = augmented['image']

        return image, torch.tensor(target, dtype=torch.long)

위의 MelonamaDataset은 image_paths 리스트, targets, augmentations를 넘겨주는 간단한 class입니다. 아이템을 얻기 위해, 해당 클래스는 PIL의 Image 모듈을 사용해서 image를 읽고, np.array로 변환하여 augmentation을 해준 뒤 target과 image를 리턴해줍니다.

우리는 train_image_paths와 val_image_paths를 얻기 위해 glob을 사용할 수 있고 그렇게 train과 val 데이터셋을 각각 만듭니다.

# psuedo code
train_image_paths = glob.glob("<path_to_train_folder>")
val_image_paths = glob.glob("<path_to_val_folder>")

sz = 256 #더 나은 AUC를 위해서 더 큰 값을 설정, 하지만 train 시간이 길어짐

train_aug = train_aug = albumentations.Compose([
    RandomCrop(sz,sz),
    ..., #원하는 augmentation
    albumentations.Normalize(always_apply=True), 
    ToTensorV2()
])

val_aug = albumentations.Compose([
    albumentations.CenterCrop(sz, sz),
    albumentations.Normalize(always_apply=True),
    ToTensorV2()
])

train_dataset = MelonamaDataset(train_image_paths, train_targets, train_aug)
val_dataset = MelonamaDataset(val_image_paths, val_targets, val_aug)

datasets이 준비되면 이제 dataloaders를 만들고 train 이미지의 sanity check를 해봅시다.

# Dataloaders
train_loader = torch.utils.data.DataLoader(
    train_dataset, batch_size=64, shuffle=True, num_workers=4)
val_loader = torch.utils.data.DataLoader(
    val_dataset, batch_size=64, shuffle=False, num_workers=4)

# visualize images
import torchvision.utils as vutils

def matplotlib_imshow(img, one_channel=False):
    fig,ax = plt.subplots(figsize=(16,8))
    ax.imshow(img.permute(1,2,0).numpy())

images= next(iter(train_loader))[0][:16]
img_grid = torchvision.utils.make_grid(images, nrow=8, normalize=True)
matplotlib_imshow(img_grid)

이제 우리의 dataloader가 완성되었어요. 아주 좋아보입니다. 우리는 분류기를 만들 준비가 끝났습니다!

Lightning Module

PL은 상용구 코드(boilerplate code)를 제거했습니다. Engineering Code와 Non-essential code를 제거함으로써써,Research code에 포커스를 맞출 수 있었어요!

 

PL의 오피셜 다큐먼트인 Quick Start 와Introduction Guide는 PL을 시작하는데 아주 훌륭한 자료입니다! 저도 그 자료들로 시작했어요.

Model and Training

 

우리가 사용할 PL모델은 다음과 같이 생겼을거에요:

class Model(LightningModule):
    def __init__(self, arch='efficientnet-b0'):
        super().__init__()
        self.base = EfficientNet.from_pretrained(arch)
        self.base._fc = nn.Linear(self.base._fc.in_features, 1)

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

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=5e-4)

    def step(self, batch):
        x, y  = batch
        y_hat = self(x)
        loss  = WeightedFocalLoss()(y_hat, y.view(-1,1).type_as(y_hat))
        return loss, y, y_hat.sigmoid()

    def training_step(self, batch, batch_nb):
        loss, y, y_hat = self.step(batch)
        return {'loss': loss}

    def validation_step(self, batch, batch_nb):
        loss, y, y_hat = self.step(batch)
        return {'loss': loss, 'y': y.detach(), 'y_hat': y_hat.detach()}

    def validation_epoch_end(self, outputs):
        avg_loss = torch.stack([x['loss'] for x in outputs]).mean()
        auc = self.get_auc(outputs)
        print(f"Epoch {self.current_epoch} | AUC:{auc}")
        return {'loss': avg_loss}

    def get_auc(self, outputs):
        y = torch.cat([x['y'] for x in outputs])
        y_hat = torch.cat([x['y_hat'] for x in outputs])
        # shift tensors to cpu
        auc = roc_auc_score(y.cpu().numpy(), 
                            y_hat.cpu().numpy()) 
        return auc

우리는 이전 포스트에서 얘기한 WeightedFocalLoss를 사용할거고, 그 이유는 데이터셋이 약 1.77 %의 긍정 클래스를 가진 불균형성때문입니다.

Model implementation compared to PyTorch

우리는 purePyTorch 처럼 __init__과 forward 메소드를 추가할 것입니다. LightningModule은 부가 기능들을 추가합니다.

 

Pure pytorch에서, training과 validation이 포함되어있는 main의 루프는 다음과 같이 생겼을겁니다:

train_dataset, valid_dataset = MelonamaDataset(...), MelonamaDatasaet(...)
train_loader, valid_loader = DataLoader(train_dataset, ...), DataLoader(valid_dataset, ...)
optimizer = ...
scheduler = ...
train_augmentations = albumentations.Compose([...])
val_aug = albumentations.Compose([...])
early_stopping = EarlyStopping(...)
model = PyTorchModel(...)
train_loss = train_one_epoch(model, optimizer, scheduler)
preds, valid_loss = evaluate(args, valid_loader, model)
report_metrics()
if early_stopping.early_stop:
    save_model_checkpoint()
    stop_training()

그리고 당연히, 우리는 train_one_epoch과 evaluate 함수를 정의할 것 입니다.

model.train()
for b_idx, data in enumerate(train_loader):
    loss = model(**data)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

evaluate 역시 아주 비슷합니다. 당신이 보고있는 것처럼, 우리는 PyTorch에서 작동하도록 많은 코드를 작성해야합니다.유연성이 뛰어나지만, 일반적으로 우리는 다양한 프로젝트에서 동일한 코드를 반복해서 재사용해야합니다.Training과 evaluate 루프는 크게 변하지 않습니다.

 

PL이 하는 일은 이 프로세스를 우리를 위해 자동화 시켜준다는 것입니다. 더 이상 상용구 코드를 작성할 필요가 없습니다.

 

Training 루프는 training_step 메소드에 있고, validation 루프는 validation_step 메소드에 들어가있습니다. 메트릭의 일반적인 리포팅은 validation_epoch_end 메소드에 있습니다. Model 클래스 안에는, training_step과 validation_step 모두 배치(batch)에서 x와 y를 가져오기 위해 step 메소드를 호출합니다. 또한 forward pass와 loss를 리턴하기 위해 foward를 호출합니다. Training이 끝나면, 우리의 validation 루프가 호출되고 epoch이 끝날 때 validation_epoch_end가 호출되서 결과가 누적되고 AUC score가 계산됩니다. 우리는 AUC가 캐글 컴피티션의 지표이기 때문에 roc_auc_score를 사용합니다.

 

짜잔 이게 다에요. 이것이 딥러닝 모델을 생성(create), 교육(train) 및 검증(validate)하는 데 필요한 모든 것입니다. 로깅을 위한 멋진 기능도 있습니다 - Wandb 및 tensorboard 그에 대한 정보는 여기에서 자세히 읽을 수 있습니다.

 

PyTorch에서 PL로 옮기는 것은 굉장히 쉽습니다. 저는 PL 소개 도큐먼트를 읽고 ISIC 모델을 PL로 몇 시간만에 구현하였습니다. 저는 PL 코드가 훨씬 조직적이고 컴팩트 하다는 것과 유연하고 실험하기에 좋다는 것을 알아냈습니다. 또한, 다른 사람들과 솔루션을 공유 할 때, 모든 사람들은 어디를보아야하는지 정확하게 알게 됩니다 - 예를 들어, training 루프는 항상 training_step에 있고, validation 루프는 항상 validation_step에 존재합니다. 다른 것들도 마찬가지입니다.

 

어떤면에서, 저는 두 라이브러리가 우리의 삶을 더 쉽게 만들어 준다는 점에서 훌륭한 fastai 라이브러리와 비교할 수있었습니다.

 

fastai와 비슷하게, PL로 모델을 train하는 것은,Trainer를 간단하게 만들고 .fit()을 호출하면 됩니다.

debug = False
gpus = torch.cuda.device_count()
trainer = Trainer(gpus=gpus, max_epochs=2, 
                  num_sanity_val_steps=1 if debug else 0)

trainer.fit(model, train_dataloader=train_loader, val_dataloaders=val_loader)

## outputs
>>  Epoch 0 | AUC:0.8667878706561116
    Epoch 1 | AUC:0.8867006574533746

짜잔 정말 이게 끝이에요. 이것이 PL에서 베이스라인 모델을 만드는 전부입니다.

Gradient Accumulation

자 여기까지 우리의 베이스라인 모델이 완성되었습니다, 이제는 gradient accumulation을 해봅시다!

trainer = Trainer(gpus=1, max_epochs=2, 
                  num_sanity_val_steps=1 if debug else 0, 
                  accumulate_grad_batches=2)

이것은 PL에 단일 매개변수를 추가하는 것 만큼 간단합니다!

아마도 일반적인 PyTorch의 워크플로우는 아래와 같을거에요:

accumulate_grad_batches=2
optimizer.zero_grad()
for b_idx, data in enumerate(train_loader):
    loss = model(**data, args=args, weights=weights)
    loss.backward()
        if (b_idx + 1) % accumulate_grad_batches == 0:
                # take optimizer every `accumulate_grad_batches` number of times
                optimizer.step()
                optimizer.zero_grad()

PL 은 이 상용구 코드를 멋지게 가져오고 리서처들이 쉽게 접근하여 gradient accumulation을 구현할 수 있도록 해줍니다. 단일 GPU에서 더 큰 배치 크기를 갖는 것이 매우 도움이 됩니다. 이것에 대해 더 읽으려면, Hugging Face의 this great article를 참조하세요!

16-bit precision training

16-bit precision은 메모리 사용량을 절반으로 줄이고 training 속도를 크게 높일 수 있습니다. 여기 이 논문은 16 bit precision에 대한 포괄적인 분석을 이야기합니다.

 

좀 더 부드럽게 소개하려면 여기에 훌륭한 자료가 있고 혼합 정밀도를 매우 잘 설명하는 fastai 문서를 참조하십시오.

더 많은 이야기는 fastai 도큐먼트에 언급되어 있습니다.

 

16-bit precision training을 추가하기 위해, 우리는 처음으로 PyTorch 1.6+ 가 필요합니다. PyTorch는 최근에 Mixed Precision Training에 대한 기본 지원만 추가했습니다.

 

가장 최근 버전의 PyTorch를 설치하기 위해선 간단히 아래의 코드를 실행하세요.

!pip install --pre torch==1.7.0.dev20200701+cu101 torchvision==0.8.0.dev20200701+cu101 -f https://download.pytorch.org/whl/nightly/cu101/torch_nightly.html

그 다음, 16-bit training을 추가하는 것은 아래처럼 간단해요:

trainer = Trainer(gpus=1, max_epochs=2, 
                  num_sanity_val_steps=1 if debug else 0, 
                  accumulate_grad_batches=2, precision=16)

만약 당신이 오래된 버전의 PyTorch로 진행하고 싶다면, 이곳을 참고하세요.

 

일반적인 PyTorch의 워크플로우에서, 16-bit precision을 위해 우리는 NVIDIA의 amp를 사용하여 training 루프를 직접 조작해서 매우 번거롭고 시간도 오래걸리는 것을 견뎌야만 합니다. PyTorch는 이제 mixed precision 및 PL에 대한 지원을 추가함으로써 구현이 매우 쉽습니다.

TPU Support

마지막으로 제가 약속했던 TPU 지원을 추가하고 TPUs에서 이 스크립트를 실행하는 것을 해볼 것입니다!

 

여기 구글의 포스트에서 TPUs를 소개하고, 여기에 멋진 블로그 포스트에서 다양한 하드웨어에 대한 비교를 해놓았습니다. TPUs는 일반적으로 V100보다 5배 빠르고 training 타임을 상당히 줄여줍니다.

 

TPU를 사용할 때, 무료 TPU를 사용하려면 Google Colab이나 Kaggle 노트북으로 스위칭해야합니다. TPUs에 대한 더 많은 정보는, 구글의 이 비디오를 보세요.

 

PL에서 TPU로 당신의 모델을 학습하는 것은 매우 간단합니다. 필요한 라이브러리를 다운로드하고 트레이너에 매개 변수를 추가하세요. :)

!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py
!python pytorch-xla-env-setup.py --version nightly --apt-packages libomp5 libopenblas-dev

trainer = Trainer(gpus=1, max_epochs=2, 
                  num_sanity_val_steps=1 if debug else 0, 
                  accumulate_grad_batches=2, precision=16, 
                  tpu_cores=8)

다음은 ISIC를 위해 pure PyTorch로 TPU를 사용하는 Abhishek Thakur의 노트북입니다. 비교해보면 PL을 사용하여 TPU를 훈련하는 것이 얼마나 쉬운 지 알게 될 것입니다.

결론

그래서 지금까지 PyTorch와 PL의 차이점을 비교할 수 있었으며, 저는 적어도 당신에게 PL을 시험하게 할만큼 설득했습니다. [여기]는 그 기술을 연습하고 PL을 사용하는 훌륭한 캐글 대회입니다! PL을 사용한 최초의 몇 가지 실험에서 저는 작업이 더 능률화되고 버그가 줄어드는 것을 발견했습니다. 다양한 배치 크기, mixed precision, 손실 함수(loss functions), 옵티 마이저 및 스케줄러로 실험하는 것이 더 쉽다는 것을 알았습니다. PL은 꼭 시도해 볼 가치가 있습니다.

Credits

읽어 주셔서 감사합니다! PyTorch Lightning을 사용해 본 결과, 이것이 실험 워크 플로에 미치는 영향에 대해서는 트위터를 통해 알려주세요. 건설적인 피드백은 언제나 환영합니다.

• 모델의 구현은 Kaggle의 멋진 노트북에서 수정 및 수정되었습니다.