Apa itu PyTorch Lightning? PyTorch Lightning adalah sebuah high-level framework berbasis PyTorch yang dirancang untuk membuat proses pelatihan deep learning menjadi lebih terstruktur, bersih, dan mudah diskalakan. Lightning memisahkan logika risetdari logika engineering (seperti training loop, logging, distribusi, dll), sehingga peneliti dan engineer bisa fokus pada model dan data, bukan pada kode boilerplate.
Mengapa Menggunakan PyTorch Lightning? Mengurangi Boilerplate: Tidak perlu lagi menulis loop training/validation/testing secara manual; Mudah untuk Diskalakan: Lightning mendukung multi-GPU, TPU, distribusi, dan mixed precision secara otomatis; Terstruktur dan Terstandar: Organisasi kode lebih rapi, dengan konvensi dan standar produksi; Kompatibel dengan PyTorch: Lightning tidak menggantikan PyTorch, melainkan membungkusnya — sehingga kamu tetap memakai PyTorch API; Kompatibel dengan Ekosistem ML: Dukungan langsung untuk wandb, tensorboard, rich logging, checkpointing, early stopping, dll.
Struktur Dasar Lightning
Berikut adalah struktur utama dalam proyek PyTorch Lightning:
-
LightningModule
Tempat mendefinisikan model, loss, optimizer, dan training logic. -
Trainer
Komponen yang menjalankan training loop. -
DataModule (opsional)
Untuk memisahkan logika pemrosesan dan loading data.
Fitur Tambahan yang Sangat Berguna
-
Callbacks: EarlyStopping, ModelCheckpoint, RichProgressBar
-
Logger: TensorBoardLogger, WandbLogger, CSVLogger
-
Tuning Tools: Learning rate finder, batch size scaler
-
Distributed Training: Multi-GPU/TPU out of the box
Saya terbiasa menggunakan pytorch secara manual maka dengan Lightning sangat terbantu. Biasanya membuat class logger sendiri, loading dan saving checkpoint juga dibuat sendiri. Maka keruwetan tersebut teratasi dengan Lightning. Saya akan modifikasi class Turunan dari LightningModule seperti berikut
import numpy as np
import pytorch_lightning as pl
import torch
from library import MyDataset2, Logger, create_model2, collate_fn, TrainingSingle
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm
class LitRCNN(pl.LightningModule):
def __init__(self, folder="lightning_logs/version_1/checkpoint/"):
super().__init__()
self.model = create_model2(num_classes= 2)
self.epoch_count = 0
self.loss_total_perbatch = list()
self.folder = folder
def forward(self, images):
return self.model(images)
def training_step(self, batch, batch_idx):
images, targets = batch
loss_dict = self.model(images, targets)
loss = sum(loss_dict.values())
self.log("train_loss", loss, on_epoch=True, prog_bar=True, logger=True)
# buat hitung loss per 1 batch
self.loss_total_perbatch.append(loss.item())
return loss
def configure_optimizers(self):
params = [p for p in self.model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.001, momentum=0.9, weight_decay=0.0005)
return optimizer
def on_train_epoch_end(self):
loss_per_iterasi = np.array(self.loss_total_perbatch).mean()
# print("epoch: ", self.current_epoch, "loss: ",loss_per_iterasi)
with open(self.folder+"log.txt", "a") as log_file:
log_file.write(f"Epoch {self.current_epoch} Loss: {loss_per_iterasi}\n")
# nol kan lagi
self.loss_total_perbatch = list()
def on_train_end(self):
print("selesai epoch!: ",self.epoch_count)
# self.log_file.close()
Gunanya untuk menyimpan tersendiri log nya karena saya butuh log hasil loss tiap iterasi! bukan per batch. Kemudian training nya sebagai berikut
%autoreload
import pytorch_lightning as pl
import torch
import numpy as np
import os
from tqdm import tqdm
from torch.utils.data import Dataset, DataLoader
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.callbacks import RichProgressBar
from pytorch_lightning.loggers import TensorBoardLogger
from library import MyDataset2, Logger, create_model2, collate_fn, TrainingSingle
from library2 import LitRCNN
dataset = MyDataset2(augmentasi = True, use_class = False, prob = 0.65, max_data = 5)
num_workers = 0
batch_size = 4
train_loader = DataLoader(dataset,
batch_size = batch_size,
shuffle = True,
num_workers = num_workers,
collate_fn = collate_fn)
version = 1
logger = TensorBoardLogger("lightning_logs", name="", version=version)
folder = "lightning_logs/version_"+str(version)+"/"
checkpoint_callback = ModelCheckpoint(
dirpath = folder+"checkpoints",
filename = "epoch={epoch}-step={step}",
save_top_k = 1, # hanya simpan 1 file (yang terbaik atau terbaru)
every_n_epochs = 1,
save_last=True,
)
model = LitRCNN(folder = folder+"checkpoints/")
file = folder+"checkpoints/last.ckpt"
trainer = pl.Trainer(max_epochs=50000,
accelerator="cpu",
devices=1,
enable_progress_bar=True,
logger=logger,
default_root_dir = folder,
callbacks=[checkpoint_callback, RichProgressBar()])
# trainer.fit(model, train_loader)
if os.path.exists(file):
trainer.fit(model, train_loader, ckpt_path=file)
else:
trainer.fit(model, train_loader)
hasil running
┏━━━┳━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━┳━━━━━━━┓
┃ ┃ Name ┃ Type ┃ Params ┃ Mode ┃
┡━━━╇━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━╇━━━━━━━┩
│ 0 │ model │ Detectorn │ 41.3 M │ train │
└───┴───────┴────────────┴────────┴───────┘
Trainable params: 41.1 M
Non-trainable params: 222 K
Total params: 41.3 M
Total estimated model params size (MB): 165
Modules in train mode: 189
Modules in eval mode: 0
Epoch 675/49999 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2/2 0:00:32 • 0:00:00 0.16it/s v_num: 1.000 train_loss_step:
0.254 train_loss_epoch: 0.221