대표적인 딥러닝 프레임워크로 pytorch, tensorflow가 있습니다. 최근에는 tensorflow보다 pytorch를 선호하는 유저가 많아지는 것 같습니다.\nPyTorch Lightning 은 PyTorch에 대한 High-level 인터페이스를 제공하는 오픈소스 Python 라이브러리입니다.
pytorch-lightning을 사용하면, 코드가 깔끔하고 간결해지며, 주어진 포맷에만 맞게 작성하면 pytorch-lightning에서 제공하는 다양한 기능을 사용할 수 있습니다.
이번 포스팅에서는 pytorch-lightning의 기본적인 사용 방법을 기록합니다.
pip install pytorch-lightning
import torch\nfrom torch import nn\nimport torch.nn.functional as F\nfrom torchvision import transforms\nfrom torchvision.datasets import MNIST\nfrom torch.utils.data import DataLoader, random_split\nimport pytorch_lightning as pl\n\n\nclass LitAutoEncoder(pl.LightningModule):\n def __init__(self):\n super().__init__()\n self.encoder = nn.Sequential(nn.Linear(28 * 28, 128), nn.ReLU(), nn.Linear(128, 3))\n self.decoder = nn.Sequential(nn.Linear(3, 128), nn.ReLU(), nn.Linear(128, 28 * 28))\n\n def forward(self, x): # prediction/inference\n embedding = self.encoder(x)\n \n return embedding\n\n def training_step(self, batch, batch_idx): # train loop, forward와 독립적으로 실행합니다.\n x, y = batch\n x = x.view(x.size(0), -1)\n z = self.encoder(x)\n x_hat = self.decoder(z)\n loss = F.mse_loss(x_hat, x)\n self.log(\"train_loss\", loss) \n # 디폴트는 TensorBoardLogger에 기록되는데, WandbLogger로만 바꿔주면 wandb를 사용할 수 있습니다. \n # loss뿐만 아니라 다른 파라미터도 wandb로 보고 싶으시면 똑같이 log를 찍어주면 됩니다.\n \n return loss\n \n def validation_step(self, batch, batch_idx):\n x, y = batch\n y_hat = self(x)\n val_loss = F.cross_entropy(y_hat, y)\n \n return val_loss\n \n def test_step(self, batch, batch_idx):\n x, y = batch\n y_hat = self(x)\n loss = F.cross_entropy(y_hat, y)\n \n return loss\n \n def configure_optimizers(self):\n optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)\n \n return optimizerdataset = MNIST(os.getcwd(), download=True, transform=transforms.ToTensor())\ntrain, val = random_split(dataset, [55000, 5000])\n\nautoencoder = LitAutoEncoder()\ntrainer = pl.Trainer()\ntrainer.fit(autoencoder, DataLoader(train), DataLoader(val))이렇게하면 학습이 시작됩니다. 하지만, 가장 기초적인 방법으로 학습을 진행한 것이기 때문에 기능들을 조금 더 알아보겠습니다.
\nclass YourLightningDataModule(pl.LightningDataModule):\n def __init__(self):\n super().__init__()\n \n \n def prepare_data(self): \n # DDP/TPU에서 모든 process가 prepare_data 메소드를 통과하지 않기 때문에 stage로 나누면 안됩니다.\n # good\n download_data()\n tokenize()\n etc()\n\n # bad\n self.split = data_split\n self.some_state = some_other_state()\n self.something = else\n\n\n def setup(self, stage: Optional[str] = None): \n # setup은 모든 process에서 호출된다. train, validate, test, and predict로 데이터를 나눈다.\n data = Load_data(...)\n self.dataset['train'] = Dataset(...)\n self.dataset['valid'] = Dataset(...)\n self.dataset['test'] = Dataset(...)\n \n def train_dataloader(self):\n return AudioDataLoader(\n dataset=self.dataset['train'],\n num_workers=4,\n batch_size=32,\n )\n\n def val_dataloader(self):\n return AudioDataLoader(\n dataset=self.dataset['valid'],\n num_workers=4,\n batch_size=32,\n )\n\n def test_dataloader(self):\n return AudioDataLoader(\n dataset=self.dataset['test'],\n num_workers=4,\n batch_size=32,\n )아래는 gpu-fp16 Trainer 예시입니다. Trainer에 대한 더 자세한 사용법은 여기를 참고하시면 좋습니다.
\ntrainer = pl.Trainer(\n precision=16, # Double precision (64), full precision (32) or half precision (16)\n accelerator=dp, # Distributed_backend (dp, ddp, etc ...)\n gpus=4, # GPU 개수\n accumulate_grad_batches=1, # Gradient를 몇 개의 배치동안 누적해서 계산할 것인지\n amp_backend=\"apex\", # mixed precision backend to use (“native” or “apex”)\n auto_select_gpus=True, # 사용가능한 GPU를 알아서 잡아준다.\n check_val_every_n_epoch=1, # 몇 개의 epoch마다 validation 할 것 인지\n gradient_clip_val=5.0, # Gradient clipping을 얼마로 할 것인지\n logger=WandbLogger(project=YOUR_PROJECT_NAME), # 로그 선택\n auto_scale_batch_size=\"binsearch\", # 메모리에 적합한 가장 큰 배치 사이즈를 찾아준다.\n max_epochs=20, # 최대 epoch 수\n)autoencoder = LitAutoEncoder()\ndata_module = YourLightningDataModule()\ntrainer = pl.Trainer(...)\n\ntrainer.fit(autoencoder, data_module)