使用 LightningDataModule 切分資料
你將完成 LightningDataModule 中的 setup 方法。正確地切分資料集可確保模型在一個子集上訓練、在另一個子集上驗證,避免過度擬合。
dataset 已經預先匯入。
本練習屬於課程
Scalable AI Models with PyTorch Lightning
練習說明
- 匯入
random_split,用來將資料集切分為訓練與驗證。 - 使用
random_split將資料集切分為訓練(80%)與驗證(20%)。
動手互動練習
試著完成這個範例程式碼,體驗一下這個練習。
# Import libraries
import lightning.pytorch as pl
from torch.utils.data import ____
class SplitDataModule(pl.LightningDataModule):
def __init__(self):
super().__init__()
self.train_data = None
self.val_data = None
def setup(self, stage=None):
# Split the dataset into training (80%) and validation (20%)
self.____, self.____ = random_split(dataset, [____, ____])