简化深度学习：PyTorch Lightning 指南
在典型的PyTorch工作流程中，我们定义一个自定义的数据集类，以帮助组织电影评论及其情感标签，供我们的模型使用。具体来说，这个类会对文本进行分词，处理不同的序列长度，并返回模型将用来学习的输入ID和标签。
·
Simplifying deep learning: A guide to PyTorch Lightning — ROCm Blogs (amd.com)
2024年2月8日 by Phillip Dang.
PyTorch Lightning 是构建在 PyTorch 之上的高级封装器。其目的是简化和抽象训练 PyTorch 模型的过程。它通过抽象掉重复的样板代码，提供了一种结构化和有组织的方法来处理机器学习（ML）任务，允许您更多地专注于模型开发和实验。PyTorch Lightning 可以开箱即用地与 AMD GPU 和 ROCm 一起使用。
有关 PyTorch Lightning 的更多信息，请参考 this article.
在这个博客中，我们在 IMDb 电影评论数据集上训练一个模型，并展示如何使用 PyTorch Lightning 简化和组织代码。我们还展示了如何使用 GPU 更快地训练模型。
您可以在 GitHub 文件夹中找到与此博客文章相关的文件GitHub folder.
前提条件
要跟随本博客，您需要具备以下软件：
接下来，确保您的系统能够识别 AMD GPU：
! rocm-smi --showproductname ================= ROCm System Management Interface ================ ========================= Product Info ============================ GPU[0] : Card series: Instinct MI210 GPU[0] : Card model: 0x0c34 GPU[0] : Card vendor: Advanced Micro Devices, Inc. [AMD/ATI] GPU[0] : Card SKU: D67301 GPU[1] : Card series: Instinct MI210 GPU[1] : Card model: 0x0c34 GPU[1] : Card vendor: Advanced Micro Devices, Inc. [AMD/ATI] GPU[1] : Card SKU: D67301 =================================================================== ===================== End of ROCm SMI Log =========================
确保 PyTorch 也能识别这些 GPU：
import torch
print(f"number of GPUs: {torch.cuda.device_count()}")
print([torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())])
number of GPUs: 2 ['AMD Radeon Graphics', 'AMD Radeon Graphics']
一旦确认系统能够识别您的设备，您就可以开始使用 PyTorch 进行典型的机器学习工作流程。这包括加载和处理数据、设置训练循环、验证循环和优化器。之后，您会看到 PyTorch Lightning 如何通过一个可扩展、易于使用的框架包装所有这些模块，使其更简便。
库
在开始之前，请确保你已经安装了所有必要的库：
pip install lightning transformers datasets torchmetrics
接下来，导入你将在本文中使用的模块：
import collections import torch import torch.nn as nn from torch.optim import AdamW from torch.utils.data import DataLoader, Dataset from datasets import load_dataset from transformers import BertTokenizer, BertModel from sklearn.metrics import accuracy_score, classification_report
数据处理
在本文中，我们的数据集是IMDb电影评论，我们的任务是分类评论是正面（1）还是负面（0）。加载数据集并查看一些示例：
# 加载IMDb数据集
imdb = load_dataset("imdb")
print(imdb)
for i in range(3):
label = imdb['train']['label'][i]
review = imdb['train']['text'][i]
print('label: ', label)
print('review:', review[:100])
print()
counts = collections.Counter(imdb['train']['label'])
print(counts)
输出结果：
DatasetDict({
train: Dataset({
features: ['text', 'label'],
num_rows: 25000
})
test: Dataset({
features: ['text', 'label'],
num_rows: 25000
})
unsupervised: Dataset({
features: ['text', 'label'],
num_rows: 50000
})
})
label: 0
review: I rented I AM CURIOUS-YELLOW from my video store because of all the controversy that surrounded it w
label: 0
review: "I Am Curious: Yellow" is a risible and pretentious steaming pile. It doesn't matter what one's poli
label: 0
review: If only to avoid making this type of film in the future. This film is interesting as an experiment b
Counter({0: 12500, 1: 12500})
我们的训练和测试数据集各包含25,000个样本，50%为正面评论，50%为负面评论。接下来，处理并对文本进行标记，并构建一个非常简单的模型来分类评论。
我们的目标不是为这个数据集构建一个超级准确的模型，而是展示一个典型的机器学习工作流程，以及如何使用PyTorch Lightning组织和简化它，使其能与我们的AMD硬件无缝运行。
构建自定义数据集类
在典型的PyTorch工作流程中，我们定义一个自定义的数据集类，以帮助组织电影评论及其情感标签，供我们的模型使用。具体来说，这个类会对文本进行分词，处理不同的序列长度，并返回模型将用来学习的输入ID和标签。
class SentimentDataset(Dataset):
def __init__(self, data, tokenizer, max_length):
self.texts = data['text']
self.labels = data['label']
self.tokenizer = tokenizer
self.max_length = max_length
def __len__(self):
return len(self.texts)
def __getitem__(self, idx):
text = self.texts[idx]
label = float(self.labels[idx])
encoding = self.tokenizer(text, return_tensors='pt', max_length=self.max_length, padding='max_length', truncation=True)
return {'input_ids': encoding['input_ids'].flatten(), 'attention_mask': encoding['attention_mask'].flatten(), 'label': torch.tensor(label)}
划分数据集
有了自定义的数据集类后，将数据分成训练集和验证集，并创建一个 DataLoader 包装器。我们通常使用 PyTorch 的 DataLoader 来支持高效的数据处理/批处理，并行化、随机打乱和采样。
# Split data into 2 sets train_data = imdb['train'] val_data = imdb['test'] # Create data set objects that handle data processing train_dataset = SentimentDataset(train_data, tokenizer, max_length) val_dataset = SentimentDataset(val_data, tokenizer, max_length) # Wrap these data set objects around DataLoader for efficient data handling and batching train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=True)
建模
接下来，创建模型。我们将使用从 Hugging Face 预训练的 Bert 模型，并对我们的文本分类任务进行微调。
我们的模型包括：
-
一个 Bert 模型，包含多个 Transformers 层。这些层执行句子中单词的自注意力上下文嵌入。
-
一个全连接的线性层，通过将输出维度从嵌入维度减少到 1，帮助在分类任务中微调模型。
class SentimentClassifier(nn.Module):
def __init__(self, pretrained_model_name='bert-base-uncased', freeze_bert=True):
super(SentimentClassifier, self).__init__()
self.bert = BertModel.from_pretrained(pretrained_model_name)
self.embedding_dim = self.bert.config.hidden_size
self.fc = nn.Linear(self.embedding_dim, 1)
if freeze_bert:
for param in self.bert.parameters():
param.requires_grad = False
def forward(self, input_ids, attention_mask):
outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
pooled_output = outputs.pooler_output
logits = self.fc(pooled_output)
return logits.squeeze(1)
初始化模型
将计算设备设置为 AMD GPU，初始化模型，创建优化器，并设置损失函数的准则。
# 设置设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 初始化模型
model = SentimentClassifier()
# 设置优化器
optimizer = AdamW(model.parameters(), lr=learning_rate)
# 设置损失函数
criterion = nn.BCEWithLogitsLoss()
不使用Lightning进行训练
这是我们在不使用Lightning的情况下，典型的训练和验证循环：
# 训练循环
model.to(device)
for epoch in range(num_epochs):
model.train()
for batch in train_dataloader:
optimizer.zero_grad()
input_ids = batch['input_ids'].to(device)
attention_mask = batch['attention_mask'].to(device)
labels = batch['label'].to(device)
preds = model(input_ids, attention_mask)
loss = criterion(preds, labels)
loss.backward()
optimizer.step()
# 评估
model.eval()
predictions = []
actual_labels = []
with torch.no_grad():
for batch in val_dataloader:
input_ids = batch['input_ids'].to(device)
attention_mask = batch['attention_mask'].to(device)
labels = batch['label'].to(device)
preds = model(input_ids, attention_mask)
predictions.extend(torch.round(preds).cpu().tolist())
actual_labels.extend(labels.cpu().tolist())
accuracy = accuracy_score(actual_labels, predictions)
print(f"Epoch [{epoch + 1}/{num_epochs}], Train Loss: {loss.item():.4f}, Val Accuracy: {accuracy:.4f}")
Epoch [1/10], Train Loss: 0.3400, Val Accuracy: 0.3130
使用PyTorch Lightning
使用PyTorch Lightning训练模型需要两个主要组件：
-
LightningDataModule: 这个模块有助于组织和封装与模型相关的所有数据逻辑。它负责准备数据、处理数据加载、预处理、将数据分割为训练/验证/测试集，以及设置数据加载器。通过将数据相关代码从模型的其他部分分离开来，它增强了代码的可读性、可重用性和可扩展性。 -
LightningModule: 这个模块封装了模型的所有方面——模型架构、优化、损失函数以及训练/验证步骤。它为这些组件提供了专门的方法，比如training_step和validation_step，定义了在训练过程中每个阶段发生的事情。
导入Lightning库以实际运行这些两个组件：
import lightning as L
使用Lightning进行数据处理
LightningDataModule 的方法通常包括：
-
__init__：初始化数据模块并设置其初始参数。这是定义属性的地方，比如批量大小 (batch size)、数据集路径、转换 (transforms) 等。同时，这里也可能会初始化在不同数据相关方法中使用的变量。 -
prepare_data：用于任何仅需要执行一次的数据相关设置，比如下载数据集或预处理原始数据。它通常在训练开始前的单个进程中运行。 -
setup：处理可能依赖于当前进程或 GPU 状态的数据相关逻辑，例如将数据集划分为训练集、验证集和测试集，或应用特定的转换。 -
train_dataloader：定义如何加载、批处理和随机打乱训练数据集。 -
val_dataloader：定义如何加载、批处理和随机打乱验证数据集.
class SentimentDataModule(L.LightningDataModule):
def __init__(self, batch_size=32, max_length=128):
super().__init__()
self.batch_size = batch_size
self.max_length = max_length
self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
def prepare_data(self):
self.dataset = load_dataset("imdb")
def setup(self, stage=None):
train_data = self.dataset['train']
val_data = self.dataset['test']
self.train_dataset = SentimentDataset(train_data, self.tokenizer, self.max_length)
self.val_dataset = SentimentDataset(val_data, self.tokenizer, self.max_length)
def train_dataloader(self):
return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True)
def val_dataloader(self):
return DataLoader(self.val_dataset, batch_size=self.batch_size, shuffle=True)
使用 Lightning 进行训练
在 PyTorch Lightning 中，核心组件是一个轻量级的 PyTorch 封装，它简化了神经网络的训练。它将 PyTorch 模型的重要部分封装到一个结构化和有组织的格式中，使得构建、训练和测试深度学习模型变得更容易。这些方法包括：
-
__init__: 定义模型的组件，如层、损失函数、优化器以及训练所需的任何其他属性。 -
forward方法: 与 PyTorch 模型类似，此方法定义神经网络的前向传播。它接受输入张量并计算输出或预测值。 -
training_step: 此方法定义单次训练迭代发生的事情。它接受一个数据批次并计算损失。它消除了显式定义批次循环、损失计算、梯度和优化器步骤的需要。 -
configure_optimizers: 定义训练期间使用的优化器，以及可选的学习率调度器。 -
validation_step: 与training_step类似，此方法定义验证迭代的计算。
class SentimentClassifier(L.LightningModule):
def __init__(self, vocab_size, embedding_dim, learning_rate=0.001) -> None:
super().__init__()
self.learning_rate = learning_rate
self.embedding = nn.Embedding(vocab_size, embedding_dim)
self.fc = nn.Linear(embedding_dim, 1)
self.criterion = nn.CrossEntropyLoss()
def forward(self, x):
embeds = self.embedding(x)
out = torch.mean(embeds, dim=1) # 在序列长度方向上平均嵌入表示
# 使用sigmoid激活函数进行二分类
out = torch.sigmoid(self.fc(out)).squeeze(1)
return out
def configure_optimizers(self):
optimizer = torch.optim.AdamW(self.parameters(), lr=self.learning_rate)
return optimizer
def training_step(self, train_batch):
x = train_batch['input_ids']
y = train_batch['label']
outputs = self(x)
loss = self.criterion(outputs, y)
self.log("train_loss", loss, prog_bar=True, on_step=False, on_epoch=True)
return loss
def validation_step(self, val_batch):
x = val_batch['input_ids']
y = val_batch['label']
outputs = self(x)
accuracy = Accuracy(task='binary').to(torch.device('cuda'))
acc = accuracy(outputs, y)
self.log('accuracy', acc, prog_bar=True, on_step=False, on_epoch=True)
return
定义了数据和模型组件后，可以使用以下代码训练模型：
model = SentimentClassifier() dm = SentimentDataModule() trainer = L.Trainer(max_epochs=num_epochs, accelerator='gpu') trainer.fit(model, dm)
GPU available: True (cuda), used: True TPU available: False, using: 0 TPU cores IPU available: False, using: 0 IPUs HPU available: False, using: 0 HPUs LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0] | Name | Type | Params ------------------------------------------------ 0 | bert | BertModel | 109 M 1 | fc | Linear | 769 2 | criterion | BCEWithLogitsLoss | 0 ------------------------------------------------ 109 M Trainable params 0 Non-trainable params 109 M Total params 437.932 Total estimated model params size (MB)
PyTorch Lightning的完整代码
以下是完整代码，你可以在notebook 或终端脚本中运行：
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from datasets import load_dataset
from transformers import BertTokenizer, BertModel
import lightning as L
from torchmetrics import Accuracy
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
class SentimentDataset(Dataset):
def __init__(self, data, tokenizer, max_length):
self.texts = data['text']
self.labels = data['label']
self.tokenizer = tokenizer
self.max_length = max_length
def __len__(self):
return len(self.texts)
def __getitem__(self, idx):
text = self.texts[idx]
label = float(self.labels[idx])
encoding = self.tokenizer(text, return_tensors='pt', max_length=self.max_length, padding='max_length', truncation=True)
return {'input_ids': encoding['input_ids'].flatten(), 'attention_mask': encoding['attention_mask'].flatten(), 'label': torch.tensor(label)}
class SentimentDataModule(L.LightningDataModule):
def __init__(self, batch_size=32, max_length=128):
super().__init__()
self.batch_size = batch_size
self.max_length = max_length
self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
def prepare_data(self):
self.dataset = load_dataset("imdb")
def setup(self, stage=None):
train_data = self.dataset['train']
val_data = self.dataset['test']
self.train_dataset = SentimentDataset(train_data, self.tokenizer, self.max_length)
self.val_dataset = SentimentDataset(val_data, self.tokenizer, self.max_length)
def train_dataloader(self):
return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True)
def val_dataloader(self):
return DataLoader(self.val_dataset, batch_size=self.batch_size, shuffle=True)
class SentimentClassifier(L.LightningModule):
def __init__(self, pretrained_model_name='bert-base-uncased', learning_rate=0.001):
super().__init__()
self.learning_rate = learning_rate
self.bert = BertModel.from_pretrained(pretrained_model_name)
self.embedding_dim = self.bert.config.hidden_size
self.fc = nn.Linear(self.embedding_dim, 1)
self.criterion = nn.BCEWithLogitsLoss()
def forward(self, input_ids, attention_mask):
outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
pooled_output = outputs.pooler_output
logits = self.fc(pooled_output)
return logits.squeeze(1)
def configure_optimizers(self):
optimizer = torch.optim.AdamW(self.parameters(), lr=self.learning_rate)
return optimizer
def training_step(self, batch):
input_ids = batch['input_ids']
attention_mask = batch['attention_mask']
labels = batch['label']
preds = self(input_ids, attention_mask)
loss = self.criterion(preds, labels)
self.log("train_loss", loss, prog_bar=True, on_step=False, on_epoch=True)
return loss
def validation_step(self, batch):
input_ids = batch['input_ids']
attention_mask = batch['attention_mask']
labels = batch['label']
preds = self(input_ids, attention_mask)
accuracy = Accuracy(task='binary').to(torch.device('cuda'))
acc = accuracy(preds, labels)
self.log('accuracy', acc, prog_bar=True, on_step=False, on_epoch=True)
return
num_epochs = 5
model = SentimentClassifier()
dm = SentimentDataModule()
trainer = L.Trainer(max_epochs=num_epochs, accelerator='gpu',limit_val_batches=0.1)
trainer.fit(model, dm)
最后，由于我们的目的是展示典型的机器学习工作流程，而不是构建最佳模型，我们鼓励你进一步实验以改进模型。我们建议尝试其他数据处理技术和模型架构，并调整超参数。
腾讯云面向开发者汇聚海量精品云计算使用和开发经验，营造开放的云计算技术生态圈。
更多推荐
23
0- 0
已为社区贡献7条内容
所有评论(0)