PyTorch Lightning — это легкая оболочка PyTorch, которая упрощает процесс создания и обучения сложных моделей глубокого обучения. Он предоставляет высокоуровневый интерфейс и лучшие практики для организации кода, обработки распределенного обучения и управления конфигурациями экспериментов. В этом сообщении блога мы рассмотрим практическое применение PyTorch Lightning и продемонстрируем, как он может ускорить ваши проекты PyTorch.
Монтаж
Прежде чем мы углубимся в примеры, давайте установим PyTorch Lightning с помощью pip, установщика пакетов Python.
pip install pytorch-lightning
Создание модели с помощью PyTorch Lightning
PyTorch Lightning предоставляет структурированный и модульный способ определения ваших моделей. Рассмотрим пример построения простой модели классификации изображений с помощью PyTorch Lightning.
import torch import torch.nn as nn from torch.optim import Adam from torch.utils.data import DataLoader from torchvision.datasets import MNIST from torchvision.transforms import ToTensor from pytorch_lightning.core import LightningModule from pytorch_lightning import Trainer class LightningClassifier(LightningModule): def __init__(self): super().__init__() self.model = nn.Sequential( nn.Linear(784, 128), nn.ReLU(), nn.Linear(128, 10), nn.LogSoftmax(dim=1) ) def forward(self, x): return self.model(x) def training_step(self, batch, batch_idx): x, y = batch logits = self(x) loss = nn.functional.nll_loss(logits, y) self.log('train_loss', loss) return loss def configure_optimizers(self): return Adam(self.parameters(), lr=0.001) # Prepare the data train_dataset = MNIST(root='./data', train=True, transform=ToTensor(), download=True) train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True) # Create the model model = LightningClassifier() # Create the trainer and train the model trainer = Trainer(max_epochs=5) trainer.fit(model, train_loader)
Распределенное обучение с PyTorch Lightning
PyTorch Lightning упрощает распределенное обучение на нескольких графических процессорах или даже на нескольких машинах. Рассмотрим пример обучения модели на нескольких GPU с помощью PyTorch Lightning.
import torch import torch.nn as nn from torch.optim import Adam from torch.utils.data import DataLoader from torchvision.datasets import MNIST from torchvision.transforms import ToTensor from pytorch_lightning.core import LightningModule from pytorch_lightning import Trainer class LightningClassifier(LightningModule): # Model definition goes here # Prepare the data train_dataset = MNIST(root='./data', train=True, transform=ToTensor(), download=True) train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True) # Create the model model = LightningClassifier() # Create the trainer and train the model on multiple GPUs trainer = Trainer(gpus=2) trainer.fit(model, train_loader)
PyTorch Lightning предоставляет интерфейс более высокого уровня для создания и обучения моделей PyTorch, что позволяет ускорить разработку и упростить эксперименты. В этом сообщении блога мы рассмотрели практическое применение PyTorch Lightning, в том числе создание модели со структурированным кодом, определение этапов обучения и настройку оптимизаторов. Мы также продемонстрировали, как использовать PyTorch Lightning для распределенного обучения на нескольких графических процессорах.
С PyTorch Lightning вы можете оптимизировать свои проекты глубокого обучения, сосредоточиться на разработке моделей и экспериментах, а также воспользоваться преимуществами эффективности и масштабируемости, которые он предлагает.
Связаться с автором: https://linktr.ee/harshita_aswani
Ссылка: