최신 딥러닝 하드웨어를 활용하려면, 하드웨어 지원 라이브러리를 사용해 딥러닝 모델을 개발해야 한다.
PyTorch Lightning은 TPU와 같은 새로운 하드웨어에서 실행되도록, 모델을 간단하게 확장할 수 있을 뿐만 아니라, 일반적인 CPU와 GPU 간 전환 프로세스를 단순화할 수 있다. 아울러, 분산 학습도 쉽게 만든다.
PyTorch Lightning 기반 그래프 신경망 벤치마크
PyTorch Geometric을 사용하여 모델을 구축하고, 내장된 데이터 세트인 "PROTEINS"을 이용해 학습해 본다.
이 데이터세트 에는 Ivanov 등 논문에 설명된 데이터 세트 정리된 버전이 포함되어 있다 . 이 데이터세트 이외에 PyTorch Geometric의 데이터 세트 모듈의 일부로 사용할 수 있다. MNIST 의 그래프 버전을 분류 하거나 선택한 그래프 분류 문제로 아래 코드에 사용된 데이터 세트를 대체할 수 있다.
그래프 딥 러닝 라이브러리인 PyTorch Geometric은 에지(인접 행렬), 노드 특성, 에지 특성 및 그래프 인덱스를 나타내는 단일 매트릭스 세트로 묶어진다.
PyTorch Lightning 및 PyTorch Geometric 가상 환경 설정
PyTorch Geometric을 Lightning과 함께 사용하고 프로젝트를 자체 가상 환경에 보관할 것이다.
virtualenv venv
source venv/bin/activate
pip install pytorch-lightning
pip install torch-scatter
pip install torch-sparse
pip install torch-cluster
pip install torch-spline-conv
pip install torch-geometric
다음으로 프로젝트 구축을 할 것이다. 이 프로젝트는 데이터셋 가져오기, 모델 정의, 교육 루프 정의, 데이터 세트 및 데이터 로더를 인스턴스화하고, 학습 루프를 호출하기 위한 루틴으로 구성된다.
라이브러리 임포트
그래프 분류를 위한 모듈을 작성할 때 가장 먼저 해야 할 일은 라이브러리 임포트이다. 다음과 같이 코딩한다.
import time
import numpy as np
import torch
from torch.nn import Dropout, Linear, ReLU
import torch_geometric
from torch_geometric.datasets import TUDataset, GNNBenchmarkDataset
from torch_geometric.loader import DataLoader
from torch_geometric.nn import GCNConv, Sequential, global_mean_pool
# this import is only used in the plain PyTorch+Geometric version
from torch.utils.tensorboard import SummaryWriter
# these imports are only used in the Lighning version
import pytorch_lightning as pl
import torch.nn.functional as F
모델 구축
PyTorch Geometric를 이용해 그래프 컨벌루션 네트워크를 구축한다. PlainGCN모델은 nn.Module 클래스에서 상속하고, PyTorch Geometric의 모델 클래스를 사용하여 순방향 학습 경로를 정의한다.
성능 비교를 위해, 파이토치 버전과 파이토치 라이트닝 버전 두개 클래스 소스를 준비한다.
class PlainGCN(torch.nn.Module):
def __init__(self, **kwargs):
super(PlainGCN, self).__init__()
self.num_features = kwargs["num_features"] \
if "num_features" in kwargs.keys() else 3
self.num_classes = kwargs["num_classes"] \
if "num_classes" in kwargs.keys() else 2
# hidden layer node features
self.hidden = 256
self.model = Sequential("x, edge_index, batch_index", [
(GCNConv(self.num_features, self.hidden),
"x, edge_index -> x1"),
(ReLU(), "x1 -> x1a"),
(Dropout(p=0.5), "x1a -> x1d"),
(GCNConv(self.hidden, self.hidden), "x1d, edge_index -> x2"),
(ReLU(), "x2 -> x2a"),
(Dropout(p=0.5), "x2a -> x2d"),
(GCNConv(self.hidden, self.hidden), "x2d, edge_index -> x3"),
(ReLU(), "x3 -> x3a"),
(Dropout(p=0.5), "x3a -> x3d"),
(GCNConv(self.hidden, self.hidden), "x3d, edge_index -> x4"),
(ReLU(), "x4 -> x4a"),
(Dropout(p=0.5), "x4a -> x4d"),
(GCNConv(self.hidden, self.hidden), "x4d, edge_index -> x5"),
(ReLU(), "x5 -> x5a"),
(Dropout(p=0.5), "x5a -> x5d"),
(global_mean_pool, "x5d, batch_index -> x6"),
(Linear(self.hidden, self.num_classes), "x6 -> x_out")])
def forward(self, graph_data):
x, edge_index, batch = graph_data.x, graph_data.edge_index,\
graph_data.batch
x_out = self.model(x, edge_index, batch)
return x_out
다음은 라이트닝 라이브러리를 사용한 코드이다.
class LightningGCN(pl.LightningModule):
def __init__(self, **kwargs):
super(LightningGCN, self).__init__()
self.num_features = kwargs["num_features"] \
if "num_features" in kwargs.keys() else 3
self.num_classes = kwargs["num_classes"] \
if "num_classes" in kwargs.keys() else 2
# hidden layer node features
self.hidden = 256
self.model = Sequential("x, edge_index, batch_index", [\
(GCNConv(self.num_features, self.hidden), \
"x, edge_index -> x1"),
(ReLU(), "x1 -> x1a"),\
(Dropout(p=0.5), "x1a -> x1d"),\
(GCNConv(self.hidden, self.hidden), "x1d, edge_index -> x2"), \
(ReLU(), "x2 -> x2a"),\
(Dropout(p=0.5), "x2a -> x2d"),\
(GCNConv(self.hidden, self.hidden), "x2d, edge_index -> x3"), \
(ReLU(), "x3 -> x3a"),\
(Dropout(p=0.5), "x3a -> x3d"),\
(GCNConv(self.hidden, self.hidden), "x3d, edge_index -> x4"), \
(ReLU(), "x4 -> x4a"),\
(Dropout(p=0.5), "x4a -> x4d"),\
(GCNConv(self.hidden, self.hidden), "x4d, edge_index -> x5"), \
(ReLU(), "x5 -> x5a"),\
(Dropout(p=0.5), "x5a -> x5d"),\
(global_mean_pool, "x5d, batch_index -> x6"),\
(Linear(self.hidden, self.num_classes), "x6 -> x_out")])
def forward(self, x, edge_index, batch_index):
x_out = self.model(x, edge_index, batch_index)
return x_out
def training_step(self, batch, batch_index):
x, edge_index = batch.x, batch.edge_index
batch_index = batch.batch
x_out = self.forward(x, edge_index, batch_index)
loss = F.cross_entropy(x_out, batch.y)
# metrics here
pred = x_out.argmax(-1)
label = batch.y
accuracy = (pred == label).sum() / pred.shape[0]
self.log("loss/train", loss)
self.log("accuracy/train", accuracy)
return loss
def validation_step(self, batch, batch_index):
x, edge_index = batch.x, batch.edge_index
batch_index = batch.batch
x_out = self.forward(x, edge_index, batch_index)
loss = F.cross_entropy(x_out, batch.y)
pred = x_out.argmax(-1)
return x_out, pred, batch.y
def validation_epoch_end(self, validation_step_outputs):
val_loss = 0.0
num_correct = 0
num_total = 0
for output, pred, labels in validation_step_outputs:
val_loss += F.cross_entropy(output, labels, reduction="sum")
num_correct += (pred == labels).sum()
num_total += pred.shape[0]
val_accuracy = num_correct / num_total
val_loss = val_loss / num_total
self.log("accuracy/val", val_accuracy)
self.log("loss/val", val_loss)
def configure_optimizers(self):
return torch.optim.Adam(self.parameters(), lr = 3e-4)
성능 평가
텐서보드와 같은 라이브러리를 사용하여 학습 진행 상황과 체크포인트를 기록할 수 있지만, 이 예제에서는 Lightning을 사용한 학습과 PyTorch만을 비교하고 있다.
def evaluate(model, test_loader, save_results=True, tag="_default", verbose=False):
# get test accuracy score
num_correct = 0.
num_total = 0.
my_device = "cuda" if torch.cuda.is_available() else "cpu"
criterion = torch.nn.CrossEntropyLoss(reduction="sum")
model.eval()
total_loss = 0
total_batches = 0
for batch in test_loader:
pred = model(batch.to(my_device))
loss = criterion(pred, batch.y.to(my_device))
num_correct += (pred.argmax(dim=1) == batch.y).sum()
num_total += pred.shape[0]
total_loss += loss.detach()
total_batches += batch.batch.max()
test_loss = total_loss / total_batches
test_accuracy = num_correct / num_total
if verbose:
print(f"accuracy = {test_accuracy:.4f}")
results = {"accuracy": test_accuracy, \
"loss": test_loss, \
"tag": tag }
return results
학습 방법
학습 루프를 정의해 보겠다. 이것을 함수로 정의하고 평가 호출을 수행한다. 학습 데이터 로더의 배치를 통해 하나의 큰 루프 내에서 진행 상황을 저장한다.
def train_model(model, train_loader, criterion, optimizer, num_epochs=1000, \
verbose=True, val_loader=None, save_tag="default_run_"):
## call validation function and print progress at each epoch end
display_every = 1 #num_epochs // 10
my_device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(my_device)
# we'll log progress to tensorboard
log_dir = f"lightning_logs/plain_model_{str(int(time.time()))[-8:]}/"
writer = SummaryWriter(log_dir=log_dir)
t0 = time.time()
for epoch in range(num_epochs):
total_loss = 0.0
batch_count = 0
for batch in train_loader:
optimizer.zero_grad()
pred = model(batch.to(my_device))
loss = criterion(pred, batch.y.to(my_device))
loss.backward()
optimizer.step()
total_loss += loss.detach()
batch_count += 1
mean_loss = total_loss / batch_count
writer.add_scalar("loss/train", mean_loss, epoch)
if epoch % display_every == 0:
train_results = evaluate(model, train_loader, \
tag=f"train_ckpt_{epoch}_", verbose=False)
train_loss = train_results["loss"]
train_accuracy = train_results["accuracy"]
if verbose:
print(f"training loss & accuracy at epoch {epoch} = "\
f"{train_loss:.4f} & {train_accuracy:.4f}")
if val_loader is not None:
val_results = evaluate(model, val_loader, \
tag=f"val_ckpt_{epoch}_", verbose=False)
val_loss = val_results["loss"]
val_accuracy = val_results["accuracy"]
if verbose:
print(f"val. loss & accuracy at epoch {epoch} = "\
f"{val_loss:.4f} & {val_accuracy:.4f}")
else:
val_loss = float("Inf")
val_acc = - float("Inf")
writer.add_scalar("loss/train_eval", train_loss, epoch)
writer.add_scalar("loss/val", val_loss, epoch)
writer.add_scalar("accuracy/train", train_accuracy, epoch)
writer.add_scalar("accuracy/val", val_accuracy, epoch)
데이터 세트, 데이터 로더 및 통합
TUDataset PROTEINs 데이터 세트, GNNBenchmarkDataset의 MNIST 그래프 버전에서 훈련할지 여부를 선택할 수 있다.
if __name__ == "__main__":
# choose the TUDataset or MNIST,
# or another graph classification problem if preferred
dataset = TUDataset(root="./tmp", name="PROTEINS")
#dataset = GNNBenchmarkDataset(root="./tmp", name="MNIST")
# shuffle dataset and get train/validation/test splits
dataset = dataset.shuffle()
num_samples = len(dataset)
batch_size = 32
num_val = num_samples // 10
val_dataset = dataset[:num_val]
test_dataset = dataset[num_val:2 * num_val]
train_dataset = dataset[2 * num_val:]
train_loader = DataLoader(train_dataset, batch_size=batch_size)
test_loader = DataLoader(test_dataset, batch_size=batch_size)
val_loader = DataLoader(val_dataset, batch_size=batch_size)
for batch in train_loader:
num_features = batch.x.shape[1]
num_classes = dataset.num_classes
# TODO
앞의 코드의 # TODO 밑에 다음 코드를 추가하고 실행해 본다.
일반적인 버전 코드는 다음과 같다.
plain_model = PlainGCN(num_features=num_features, num_classes=num_classes)
lr = 1e-5
num_epochs = 2500
criterion = torch.nn.CrossEntropyLoss(reduction="sum")
optimizer = torch.optim.Adam(plain_model.parameters(), lr=lr)
train_model(plain_model, train_loader, criterion, optimizer,\
num_epochs=num_epochs, verbose=True, \
val_loader=val_loader)
Lightning 버전에서는 클래스를 사용하여 Trainer을 호출하고 학습을 조정한다.
lightning_model = LightningGCN(num_features=num_features, num_classes=num_classes)
num_epochs = 2500
val_check_interval = len(train_loader)
trainer = pl.Trainer(max_epochs = num_epochs, \
val_check_interval=val_check_interval, gpus=[0])
trainer.fit(lightning_model, train_loader, val_loader)
코드 실행 결과는 다음과 같다.
테스트 결과는 다음과 같다.
마무리
PyTorch Lightning에서 Trainer 개체는 GPU, TPU 및 IPU와 같은 다양한 하드웨어 가속기에 대한 할당을 처리한다. 이를 통해 Google의 TPU 또는 Graphcore의 IPU와 같은 ASIC 하드웨어 가속기로 확장하는 것도 훨씬 쉽다.
Lightning은 각각의 새 프로젝트에 대한 상용구 코드의 양을 줄이고 하드웨어 가속을 쉽게 적용하는 좋은 방법이다.
