Gorio Tech Blog search

Pytorch Geometric custom graph convolutional layer 생성하기

05 Sep 2021 | Machine_Learning PyTorch

목차

Pytorch Geometric의 MessagePassing class에 대한 간략한 설명을 참고하고 싶다면 이 글을 확인하길 바란다.

본 글에서는 MessagePassing class를 상속받아 직접 Graph Convolutional Layer를 만드는 법에 대해서 다루도록 하겠다. 그 대상은 RGCN: Relational GCN이다. RGCN에 대한 간략한 설명을 참고하고 싶다면 이 곳을 확인해도 좋고 RGCN source code를 확인해도 좋다.

Custom GCN layer 생성하기: RGCN

본 포스팅에서는 원본 source code의 형식을 대부분 보존하면서도 간단한 설명을 위해 필수적인 부분만 선별하여 설명의 대상으로 삼도록 할 것이다.

먼저 필요한 library를 불러오고 parameter를 초기화하기 위한 함수를 선언한다.

import math
from typing import Optional, Union, Tuple

from torch.nn import Parameter

from torch_geometric.typing import OptTensor, Adj
from torch_geometric.nn.conv import MessagePassing

def glorot(tensor):
    if tensor is not None:
        stdv = math.sqrt(6.0 / (tensor.size(-2) + tensor.size(-1)))
        tensor.data.uniform_(-stdv, stdv)

def zeros(tensor):
    if tensor is not None:
        tensor.data.fill_(0)

RGCN에서는 regularization 방법으로 2가지를 제시하고 있는데 본 포스팅에서는 자주 사용되는 basis-decomposition 방법을 기본으로 하여 진행하도록 하겠다.

예를 들기 위해 적합한 데이터를 생각해보자. (참고로 아래 setting은 MUTAG 데이터셋을 불러온 것이다. 아래 코드를 통해 다운로드 받을 수 있다.)

dataset = 'MUTAG'

path = os.path.join(os.getcwd(), 'data', 'Entities')
dataset = Entities(path, dataset)
data = dataset[0]
구분 설명
edge_index (2, 148454)
edge_type (148454), 종류는 46개
num_nodes 23606
x node features는 주어지지 않음

그렇다면 이 layer의 목적은 23606개의 node에 대하여, 46종류의 relation을 갖는 edges를 통해 message passing을 진행하는 것이다. MUTAG 데이터셋에는 node features는 존재하지 않지만, data.x = torch.rand((23606, 100)) 코드를 통해 가상의 데이터를 만들어서 임시로 연산 과정을 살펴볼 수 있을 것이다.

먼저 반드시 필요한 사항에 대해 정의해보자.

class RelationalGCNConv(MessagePassing):
    def __init__(self,
                 in_channels: int,
                 out_channels: int,
                 num_relations: int,
                 num_bases: Optional[int]=None,
                 aggr: str='mean',
                 **kwargs):
        super(RelationalGCNConv, self).__init__(aggr=aggr, node_dim=0)

        # aggr, node_dim은 알아서 self의 attribute로 등록해준다. (MessagePassing)
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.num_relations = num_relations
        self.num_bases = num_bases

        # 원본 코드에서 in_channels가 tuple인 이유는
        # src/tgt node가 다른 type인 bipartite graph 구조도 지원하기 위함이다.
        # 예시를 위해 integer로 변경한다.
        self.in_channels = in_channels

RGCN은 Full-batch training을 기본으로 하고 있다. 이 때문에 node_dim을 조정해 주어야 한다.

다음으로는 Weight Parameter를 정의해주자. 이해를 쉽게 하기 위해서 원본 코드에서 변수 명을 수정하였다.

        # Define Weights
        if num_bases is not None:
            self.V = Parameter(
                data=Tensor(num_bases, in_channels, out_channels))
            self.a = Parameter(Tensor(num_relations, num_bases))
        else:
            self.V = Parameter(
                Tensor(num_relations, in_channels, out_channels))
            # dummy parameter
            self.register_parameter(name='a', param=None)

        self.root = Parameter(Tensor(in_channels, out_channels))
        self.bias = Parameter(Tensor(out_channels))

        self.reset_parameters()

basis-decomposition의 식을 다시 확인해보자.

\[W_r^l = \Sigma_{b=1}^B a_{rb}^l V_b^l\]

reset_parameters 메서드는 아래와 같다.

    def reset_parameters(self):
        glorot(self.V)
        glorot(self.a)
        glorot(self.root)
        zeros(self.bias)

이제 본격적으로 forward 함수를 구현할 차례이다. 사실 RGCN의 경우 특별히 변형을 하지 않느다면, 특정 relation 안에서의 연산은 일반적인 GCN과 동일하다. 따라서 기본적 세팅에서는 message, aggregate, update 메서드를 override할 필요가 없다.

다음은 forward 함수의 윗 부분이다.

    def forward(self,
                x: OptTensor,
                edge_index: Adj,
                edge_type: OptTensor=None):

        x_l = x
        # node feature가 주어지지 않는다면
        # embedding weight(V) lookup을 위해 아래와 같이 세팅한다.
        if x_l is None:
            x_l = torch.arange(self.in_channels, device=self.V.device)

        x_r = x_l
        size = (x_l.size(0), x_r.size(0))

        # output = (num_nodes, out_channels)
        out = torch.zeros(x_r.size(0), self.out_channels, device=x_r.device)

num_bases 인자가 주어진다면 아래와 같이 weight를 다시 계산해준다.

        V = self.V
        if self.num_bases is not None:
            V = torch.einsum("rb,bio->rio", self.a, V)

자 이제 각 relation 별로 propagate를 진행해주면 된다. 앞서 언급하였듯이 특정 relation 내에서의 연산은 일반적인 GCN과 다를 것이 없다. 참고로 아래와 같이 계산하면 속도 측면에서 매우 불리한데, 이를 개선한 FastRGCNConv layer가 존재하니 참고하면 좋을 것이다. 다만 이 layer의 경우 메모리를 크게 사용하므로 본격적인 사용에 앞서 점검이 필요할 것이다.

        # propagate given relations
        for i in range(self.num_relations):
            # 특정 edge_type에 맞는 edge_index를 선별한다.
            selected_edge_index = masked_edge_index(edge_index, edge_type == i)

            # node_features가 주어지지 않는다면
            if x_l.dtype == torch.long:
                out += self.propagate(selected_edge_index, x=V[i, x_l], size=size)

            # node_features가 주어진다면
            else:
                h = self.propagate(selected_edge_index, x=x_l, size=size)
                out += (h @ V[i])

        out += self.root[x_r] if x_r.dtype == torch.long else x_r @ self.root
        out += self.bias
        return out

masked_edge_index 함수는 아래와 같다.

from torch_sparse import masked_select_nnz

def masked_edge_index(edge_index, edge_mask):
    """
    :param edge_index: (2, num_edges)
    :param edge_mask: (num_edges) -- source node 기준임
    :return: masked edge_index (edge_mask에 해당하는 Tensor만 가져옴)
    """
    if isinstance(edge_index, Tensor):
        return edge_index[:, edge_mask]
    else:
        # if edge_index == SparseTensor
        return masked_select_nnz(edge_index, edge_mask, layout='coo')

여기까지 진행했다면 custom gcn layer 구현은 끝난 것이다. 아래와 같이 사용하면 된다.

data = data.to(device)

model = RelationalGCNConv(
    in_channels=in_channels, out_channels=out_channels,
    num_relations=num_relations, num_bases=num_bases).to(device)

print(get_num_params(model))

out = model(x=data.x, edge_index=data.edge_index, edge_type=data.edge_type)

Related Posts