MLflow를 활용한 GCN 모델 예제 / 25.02.04

MLflow를 설치했으니 간단한 예제 코드로 사용법을 익혀보자.

석사 기간 동안 GNN에 대해 많이 다루어봤기 때문에 사용할 모델은 GCN으로 선택했다.

 

먼저, 사용할 예제 코드이다.

해당 코드는 Cora dataset을 활용하여 GCN 모델을 학습한 후, node의 class를 분류하는 task를 수행한다.

 

데이터 정보는 아래 github에서 확인할 수 있다.

https://github.com/Cow-Kite/GNN_Project/tree/main/Dataset

GNN_Project/Dataset at main · Cow-Kite/GNN_Project

 

import time
import torch
import torch_geometric.transforms as T
import torch.nn as nn
from torch_geometric.nn import GCNConv
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid


class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(dataset.num_node_features, 16)
        self.conv2 = GCNConv(16, dataset.num_classes)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        output = self.conv2(x, edge_index)
        return output
    

def train_node_classifier(model, graph, optimizer, criterion, n_epochs=200):

    start_time = time.time()

    for epoch in range(1, n_epochs + 1):
        model.train() 
        optimizer.zero_grad() 
        out = model(graph) 
        loss = criterion(out[graph.train_mask], graph.y[graph.train_mask]) 
        loss.backward() 
        optimizer.step() 

        pred = out.argmax(dim=1) 
        acc = eval_node_classifier(model, graph, graph.val_mask)

        if epoch % 10 == 0:
            print(f'Epoch: {epoch:03d}, Train Loss: {loss:.3f}, Val Acc: {acc:.3f}')

    total_time = time.time() - start_time
    print(f"Training completed in {total_time:.2f} seconds.")

    test_acc = eval_node_classifier(model, graph, graph.test_mask)

    print(f"Test Accuracy: {test_acc:.3f}")
    return model


def eval_node_classifier(model, graph, mask):
    model.eval() 
    pred = model(graph).argmax(dim=1)
    correct = (pred[mask] == graph.y[mask]).sum()
    
    acc = int(correct) / int(mask.sum())
    return acc


if __name__=="__main__":
    device = "cpu"

    dataset = Planetoid(root='./', name='Cora')
    graph = dataset[0]
    split = T.RandomNodeSplit(num_val=0.1, num_test=0.2)
    graph = split(graph)

    gcn = GCN().to(device)
    optimizer_gcn = torch.optim.Adam(gcn.parameters(), lr=0.01, weight_decay=5e-4)
    criterion = nn.CrossEntropyLoss()

    gcn = train_node_classifier(gcn, graph, optimizer_gcn, criterion)

 

이제 이 코드에 mlflow를 활용하여 하이퍼파라미터, 손실, 정확도 등을 자동으로 기록해보자.

▶ mlflow, uuid import

먼저, 필요한 패키지를 import 한다.

import time
import torch
import torch_geometric.transforms as T
import torch.nn as nn
from torch_geometric.nn import GCNConv
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid
# mlflow
import mlflow
import mlflow.pytorch
import uuid

---

▶ main 함수

그 다음, main 함수에서 mlflow 설정을 진행한다.

if __name__=="__main__":
    # set mlflow
    study_name="GCN-tutorial" # experiment 이름
    mlflow.set_tracking_uri("http://0.0.0.0:5001") # mlflow 서버 URL
    mlflow.set_experiment(study_name)

    device = "cpu"

    dataset = Planetoid(root='./', name='Cora')
    graph = dataset[0]
    split = T.RandomNodeSplit(num_val=0.1, num_test=0.2)
    graph = split(graph)

    gcn = GCN().to(device)
    optimizer_gcn = torch.optim.Adam(gcn.parameters(), lr=0.01, weight_decay=5e-4)
    criterion = nn.CrossEntropyLoss()

    gcn = train_node_classifier(gcn, graph, optimizer_gcn, criterion)

---

▶ run_name 이름 지정

실험의 이름은 main 함수에서 GCN-tutorial이라고 지정했다. 이제 이 코드를 실행할 때마다, GCN-tutorial에 저장된다.

이때, 실행할 때마다 저장되는 이름이 바로 run_name이다.

 

그러나, run_name이 고정 이름이면 실험이 다 똑같은 이름으로 저장된다. 

따라서, uuid를 사용하여 이름에 랜덤값을 추가해 각 실험이 구별되도록 한다.

run_name은 GCN 모델위에 선언했다.

 

run_name = f"GCN_Run_{uuid.uuid4().hex[:8]}"

class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(dataset.num_node_features, 16)
        self.conv2 = GCNConv(16, dataset.num_classes)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        output = self.conv2(x, edge_index)
        return output

---

▶ train_node_classifier 함수

이제 이 함수에서 모델의 파라미터, loss, acc 등을 기록할 수 있게 코드를 추가한다.

mlflow.start_run 함수를 사용하여 새로운 run을 실행하고 기록을 시작한다.

파라미터를 기록하는 함수는 log_param이고, metric을 기록하는 함수는 log_metric이다.

def train_node_classifier(model, graph, optimizer, criterion, n_epochs=200):
    with mlflow.start_run(run_name=run_name):
        mlflow.log_param("epochs", n_epochs)
        mlflow.log_param("learning_rate", optimizer.param_groups[0]['lr'])
        mlflow.log_param("weight_decay", optimizer.param_groups[0]['weight_decay'])

        start_time = time.time()

        for epoch in range(1, n_epochs + 1):
            model.train() 
            optimizer.zero_grad() 
            out = model(graph) 
            loss = criterion(out[graph.train_mask], graph.y[graph.train_mask]) 
            loss.backward() 
            optimizer.step() 

            pred = out.argmax(dim=1) 
            acc = eval_node_classifier(model, graph, graph.val_mask)

            mlflow.log_metric("train_loss", loss.item(), step=epoch)
            mlflow.log_metric("val_acc", acc, step=epoch)

            if epoch % 10 == 0:
                print(f'Epoch: {epoch:03d}, Train Loss: {loss:.3f}, Val Acc: {acc:.3f}')

        total_time = time.time() - start_time
        mlflow.log_metric("training_time", total_time)
        print(f"Training completed in {total_time:.2f} seconds.")

        mlflow.pytorch.log_model(model, "trained_gcn")

        test_acc = eval_node_classifier(model, graph, graph.test_mask)
        mlflow.log_metric("test_acc", test_acc) 

        print(f"Test Accuracy: {test_acc:.3f}")
        return model

---

▶ MLflow Web 

마지막으로, 수정한 코드를 실행한 뒤 mlflow web으로 접속해보자.

 

• 실험 목록에 GCN-tutorial이 생성되었다.

• GCN-tutorial를 클릭하면, 실행한 run 목록을 볼 수 있다.

• run을 클릭하면, 파라미터, 메트릭 등 실험과 관련된 다양한 정보를 확인할 수 있다.

---

▶ 마치며...

간단하게 mlflow를 사용해보았다. 

 

실제 석사 기간 동안에는 파라미터, 학습 시간, loss, 정확도 등을 엑셀에 매번 작성했었는데,

코드 몇 줄을 추가하여 수고를 덜 수 있는 부분이 참 편하다고 생각한다.

 

다음에는 minio와 연동해서 사용하는 실습을 할 예정이다.