開始使用免費開始

使用 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, [____, ____])
編輯並執行程式碼