MLOps E2E - 4. Logging : wandb
이번 단계는 kubeflow training을 모니터링하고 log를 기록하는 방법이다.
쿠버네티스의 리소스 모니터링은 주로 프로메테우스와 그라파나를 쓰는 것으로 알고 있지만, 이번 연습 단계에선 간단하게 logging과 hyperparameter tuning이 가능한 wandb를 사용하려고 한다.
프로메테우스나 그라파나는 나중에 필요할 때 천천히 추가하도록 해야겠다.
이번 글은 wandb를 training code에 합치는 것이다.
1. 물론 가장 먼저 사용하는 docker image에 wandb가 설치되어 있어야 한다.
2. 그 다음 사용할 wandb id에서 API KEY를 획득해야 한다.
(wandb 계정이 없다면 가입을 먼저 한다.)
wandb.ai 를 접속한 뒤 우측 위에 계정을 누른 뒤 setting를 들어간다.
그리고 내리다 보면 API keys 항목에 키 하나가 있다.
글씨를 누르면 복사가 알아서 된다.
Key 값을 코드에서 사용해야 하는데 바로 써도 되고 환경변수에 옮겨서 써도 된다.
나는 환경변수에 등록했다.
3. 그러고 나서 pipeline 코드에서 train component에 환경변수로 API key값을 넣어준다.
이 때 환경변수의 이름은 반드시 "WANDB_API_KEY" 이어야 한다.
기존의 pipeline.py에서 수정된 코드는 다음과 같다.
(수정된 부분 : "###########################"으로 둘러 쌓인 부분)
pipeline.py
###########################
from kubernetes.client.models import V1EnvVar
###########################
...
@dsl.pipeline(
name="Pytorch Lightning Training pipeline", description="Cifar 10 dataset pipeline"
)
def pl_pipeline():
...
train_task = (
train_op(
dataset_path=data_task.outputs["output_data"],
model=MODEL,
gpus=GPUS,
max_epochs=MAX_EPOCHS,
num_classes=NUM_CLASSES,
train_batch_size=TRAIN_BATCH_SIZE,
train_num_workers=TRAIN_NUM_WORKERS,
val_batch_size=VAL_BATCH_SIZE,
val_num_workers=VAL_NUM_WORKERS,
lr=LR,
)
.after(data_task)
.set_display_name("Training")
.add_pvolumes({"/dev/shm": shm_volume})
###########################
.add_env_variable(
V1EnvVar(name="WANDB_API_KEY", value=os.environ["WANDB_API_KEY"])
)
###########################
).set_gpu_limit(1)다음과 같이 add_env_variable 함수를 통해 WANDB_API_KEY 값을 환경변수로 추가한다.
4. 그 다음 pl_train.py 코드에서 wandb를 실행시킨다.
pytorch-lightning에서 wandb를 실행시킬 땐 pytorch_lightning.loggers.WandbLogger 클래스를 이용해야 한다.
수정된 pl_train.py는 다음과 같다.
(수정된 부분 : "###########################"으로 둘러 쌓인 부분)
pl_train.py
###########################
from pytorch_lightning.loggers import WandbLogger
import wandb
from Training.pl_logger import Logger
###########################
...
if __name__ == "__main__":
...
datamodule = CIFAR10DataModule(**datamodule_args)
datamodule.setup()
samples = next(iter(datamodule.test_dataloader()))
###########################
wandb_logger = WandbLogger(
project=args.wandb_project,
name=args.wandb_name,
log_model=False,
save_dir=log_dir,
)
###########################
# Initiating the training process
trainer = Trainer(
###########################
logger=wandb_logger, # W&B integration
###########################
log_every_n_steps=args.log_every_n_steps, # set the logging frequency
gpus=args.gpus, # use all GPUs
max_epochs=args.max_epochs, # number of epochs
deterministic=True, # keep it deterministic
enable_checkpointing=True,
callbacks=[
###########################
Logger(samples),
###########################
early_stopping,
checkpoint_callback,
], # see Callbacks section
precision=16,
check_val_every_n_epoch=args.check_val_every_n_epoch,
strategy="ddp",
auto_scale_batch_size=True,
)
model_config = {
"model_name": args.model,
"num_classes": args.num_classes,
"lr": args.lr,
"width": 32,
"height": 32,
}
model = Classifier(**model_config)
trainer.fit(model, datamodule)
trainer.test(datamodule=datamodule)
###########################
wandb.finish()
###########################WandbLogger 클래스를 먼저 선언한다.
그 다음 Trainer의 logger 인자에 선언한 WandbLogger를 입력하면 된다.
그리고 마지막에 wandb.finish()를 통해 종료를 알린다.
5. validation epoch 마다 추가적인 로그를 기록하였다.
Logger 클래스는 새롭게 만든 클래스며 Training/pl_logger.py 파일에 있다.
callbacks에 새롭게 들어간 Logger(samples)는 중간중간 validation step마다 이미지 결과를 기록하는 클래스이다.
Training/pl_logger.py
import torch
import pytorch_lightning as pl
import wandb
class Logger(pl.Callback):
def __init__(self, val_samples, num_samples=32):
super().__init__()
self.val_imgs, self.val_labels = val_samples
self.val_imgs = self.val_imgs[:num_samples]
self.val_labels = self.val_labels[:num_samples]
def on_validation_epoch_end(self, trainer, pl_module):
val_imgs = self.val_imgs.to(device=pl_module.device)
logits = pl_module(val_imgs)
preds = torch.argmax(logits, 1)
trainer.logger.experiment.log(
{
"examples": [
wandb.Image(x, caption=f"Pred:{pred}, Label:{y}")
for x, pred, y in zip(val_imgs, preds, self.val_labels)
],
}
)validation epoch이 끝날 때 마다 wandb에 image sample 32개의 결과값을 기록한다.
최종 결과
Kubeflow Dashboard
Weight and Bias
이 다음 글은 minio에 연결하여 model과 log를 minio에 저장하는 것이다.