[Torch Uncertainty] Tutorial: Training a LeNet with Monte-Carlo Dropout

`Torch Uncertainty` 공식문서의 `Tutorials`를 번역한 내용입니다.

---

Training a LeNet with Monte-Carlo Dropout

이 튜토리얼에서는 Monte-Carlo Dropout(MC Dropout)이라는 계산 효율적인 베이지안 근사 방법을 사용하여 MNIST 데이터셋에서 LeNet 분류기를 훈련합니다. 예측 평균과 불확실성(분산)을 추정하기 위해, 드롭아웃 레이어가 활성화된 상태로 여러 번 네트워크를 forward pass시킵니다.

 

Monte-Carlo Dropout에 대한 더 자세한 정보는 다음 문헌을 참조하십시오:

• "Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning" (ICML 2016)
• "What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?" (NeurIPS 2017)

Training a LeNet with MC Dropout using TorchUncertainty models and PyTorch Lightning

이 부분에서는 TU(TorchUncertainty)에 이미 구현된 모델과 루틴을 기반으로 드롭아웃 레이어가 포함된 LeNet을 훈련합니다.

1. Loading the utilities

먼저 TorchUncertainty에서 다음 유틸리티를 로드해야 합니다:

• TorchUncertainty 유틸리티에서 `TUTrainer`
• 데이터로더를 관리하는 데이터 모듈: `torch_uncertainty.datamodules에서 MNISTDataModule`
• 모델: `torch_uncertainty.models`에서 `lenet`
• MC Dropout 래퍼: `torch_uncertainty.models.wrappers`에서 `mc_dropout`
• TorchUncertainty 루틴에서의 분류 훈련 및 평가 루틴
• `torch_uncertainty.optim_recipes` 모듈에서 최적화 레시피

또한 `torch.nn` 내의 신경망 유틸리티도 임포트해야 합니다.

from pathlib import Path

from torch_uncertainty.utils import TUTrainer
from torch import nn

from torch_uncertainty.datamodules import MNISTDataModule
from torch_uncertainty.models.lenet import lenet
from torch_uncertainty.models import mc_dropout
from torch_uncertainty.optim_recipes import optim_cifar10_resnet18
from torch_uncertainty.routines import ClassificationRoutine

2. Defining the Model and the Trainer

다음 단계에서는 먼저 트레이너를 생성하고, MNIST 데이터셋, 데이터로더 및 변환을 처리하는 데이터 모듈을 인스턴스화합니다. 그런 다음, `torch_uncertainty.models`에서 제공하는 청사진(blueprint)을 사용해 모델을 생성하고, 이를 `mc_dropout`으로 래핑합니다. `mc_dropout` 래퍼를 사용하려면 드롭아웃 모듈을 함수형(functional)으로 사용하지 말고 모듈형으로 사용해야 하며, 반드시 `__init__` 메서드 내에서 인스턴스화해야 합니다.

trainer = TUTrainer(accelerator="cpu", max_epochs=2, enable_progress_bar=False)

# datamodule
root = Path("data")
datamodule = MNISTDataModule(root=root, batch_size=128)


model = lenet(
    in_channels=datamodule.num_channels,
    num_classes=datamodule.num_classes,
    dropout_rate=0.4,
)

mc_model = mc_dropout(model, num_estimators=16, last_layer=False)

3. The Loss and the Training Routine

이 문제는 분류 문제이며, 우리는 NLL(negative log likelihood)로서 `CrossEntropyLoss`를 사용합니다. 훈련 루틴은 `torch_uncertainty.routines`에 있는 분류 훈련 루틴을 사용하여 정의합니다. 여기서 클래스의 수, 최적화 레시피를 제공하고, 평가 시 모델이 앙상블로 작동하도록 루틴에 설정합니다.

routine = ClassificationRoutine(
    num_classes=datamodule.num_classes,
    model=mc_model,
    loss=nn.CrossEntropyLoss(),
    optim_recipe=optim_cifar10_resnet18(mc_model),
    is_ensemble=True,
)

4. Gathering Everything and Training the Model

이제 트레이너를 사용하여 모델을 훈련할 수 있습니다. 트레이너의 `fit` 및 `test` 메서드에 훈련 루틴과 데이터 모듈을 전달하면 됩니다. 이 과정에서 트레이너는 자동으로 몇 가지 불확실성 지표를 평가하며, 이러한 지표들은 아래 표에서 확인할 수 있습니다.

trainer.fit(model=routine, datamodule=datamodule)
results = trainer.test(model=routine, datamodule=datamodule)

5. Testing the Model

이제 모델이 훈련되었으므로 MNIST에서 테스트해 보겠습니다. 평가 시 드롭아웃을 활성화하고 여러 번(여기서는 16번) 예측을 얻기 위해 `.eval()` 메서드를 호출하는 것을 잊지 마세요.

import matplotlib.pyplot as plt
import numpy as np
import torch
import torchvision


def imshow(img):
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.axis("off")
    plt.tight_layout()
    plt.show()


dataiter = iter(datamodule.val_dataloader())
images, labels = next(dataiter)

# print images
imshow(torchvision.utils.make_grid(images[:6, ...], padding=0))
print("Ground truth labels: ", " ".join(f"{labels[j]}" for j in range(6)))

routine.eval()
logits = routine(images).reshape(16, 128, 10)

probs = torch.nn.functional.softmax(logits, dim=-1)


for j in range(6):
    values, predicted = torch.max(probs[:, j], 1)
    print(
        f"MC-Dropout predictions for the image {j+1}: ",
        " ".join([str(image_id.item()) for image_id in predicted]),
    )

Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).
Ground truth labels:  7 2 1 0 4 1
MC-Dropout predictions for the image 1:  7 7 7 7 7 7 8 7 7 7 9 7 7 7 7 7
MC-Dropout predictions for the image 2:  0 2 2 0 2 0 2 2 2 0 0 2 2 2 2 2
MC-Dropout predictions for the image 3:  1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1
MC-Dropout predictions for the image 4:  0 0 0 0 0 7 0 9 0 0 0 0 0 0 7 0
MC-Dropout predictions for the image 5:  4 4 4 4 4 4 4 4 2 9 4 4 4 4 4 4
MC-Dropout predictions for the image 6:  1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1

 

대부분의 경우, 드롭아웃을 사용한 사후 분포 근사의 샘플들 사이에 의견의 차이가 발생하는 것을 관찰할 수 있습니다.

참조

https://torch-uncertainty.github.io/auto_tutorials/tutorial_mc_dropout.html#sphx-glr-auto-tutorials-tutorial-mc-dropout-py

Training a LeNet with Monte-Carlo Dropout — TorchUncertainty 0.2.1.post0 documentation