PyTorch Geometric(PYG)-实现小批量data类中inc与__cat_dim__的含义与作用

最新推荐文章于 2024-06-30 17:00:07 发布

点云哥

最新推荐文章于 2024-06-30 17:00:07 发布

阅读量2.6k

点赞数 6

CC 4.0 BY-SA版权

本文链接：https://blog.csdn.net/qq_41795143/article/details/114281387

深度学习专栏收录该内容

6 篇文章

订阅专栏

本文解析了PyTorch Geometric (PYG) 中__inc__与__cat_dim__的作用及实现方法，通过实例展示了如何自定义这两个函数来实现特定的数据拼接需求。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

PYG中实现小批量data类中inc与__cat_dim__的含义与作用

1.作用

此两个函数出现在pytorch geometric实现批量操作时，batch集行为的自定义修改方法，两种方法都是为了解决多个数据之间的拼接问题。

2.直观图解

官方初始定义，均对某一属性值进行判定

def __inc__(self, key, value):
    if 'index' in key or 'face' in key:
        return self.num_nodes
    else:
        return 0

def __cat_dim__(self, key, value):
    if 'index' in key or 'face' in key:
        return 1
    else:
        return 0

返回值的具体含义见图：

__ inc __

即__ inc __返回值表示相应矩阵错位步数，一般用于边的邻接矩阵：
__ cat_dim __

即__ cat_dim __返回值表示相应矩阵拼接的维度。按行或列拼接，一般用于节点或者结果矩阵的拼接

3.官方示例

官方默认batch拼接形式如下在这里插入图片描述
A为邻接矩阵，X为节点矩阵，Y为结果矩阵

此时对比官方示例，一目了然：

成对数据结构

class PairData(Data):
    def __init__(self, edge_index_s, x_s, edge_index_t, x_t):
        super(PairData, self).__init__()
        self.edge_index_s = edge_index_s
        self.x_s = x_s
        self.edge_index_t = edge_index_t
        self.x_t = x_t

批量时，要分成两个数据集来拼接，所以需要自定义边矩阵的拼接

def __inc__(self, key, value):
    if key == 'edge_index_s':
        return self.x_s.size(0)
    if key == 'edge_index_t':
        return self.x_t.size(0)
    else:
        return super().__inc__(key, value)

测试

edge_index_s = torch.tensor([
    [0, 0, 0, 0],
    [1, 2, 3, 4],
])
x_s = torch.randn(5, 16)  # 5 nodes.
edge_index_t = torch.tensor([
    [0, 0, 0],
    [1, 2, 3],
])
x_t = torch.randn(4, 16)  # 4 nodes.

data = PairData(edge_index_s, x_s, edge_index_t, x_t)
data_list = [data, data]
loader = DataLoader(data_list, batch_size=2)
batch = next(iter(loader))

print(batch)
>>> Batch(edge_index_s=[2, 8], x_s=[10, 16],
          edge_index_t=[2, 6], x_t=[8, 16])

print(batch.edge_index_s)
>>> tensor([[0, 0, 0, 0, 5, 5, 5, 5],
            [1, 2, 3, 4, 6, 7, 8, 9]])

print(batch.edge_index_t)
>>> tensor([[0, 0, 0, 4, 4, 4],
            [1, 2, 3, 5, 6, 7]])