跟着官方文档学DGL框架第十三天——大型图上的随机训练之边分类

参考资料

  1. https://docs.dgl.ai/en/latest/guide/minibatch-edge.html#guide-minibatch-edge-classification-sampler
  2. https://docs.dgl.ai/en/latest/api/python/dgl.dataloading.html#dgl.dataloading.pytorch.EdgeDataLoader

概述

边分类的随机训练与节点分类的随机训练没有太大区别,因为边分类的任务第一步也需要获得节点的表示。而且边的采样与节点的采样本就差不多,邻居采样器仍可以使用节点分类时用的邻居采样器,只是需要与dgl.dataloading.EdgeDataLoader()配合使用。

随机训练之同构图上的边分类

加载数据

依然使用“跟着官方文档学DGL框架第八天”中定义的DGLDataset类型的数据集。随机为每条边打上了标签,并随机选择了100条边作为训练集。

def build_karate_club_graph():
    # All 78 edges are stored in two numpy arrays. One for source endpoints
    # while the other for destination endpoints.
    src = np.array([1, 2, 2, 3, 3, 3, 4, 5, 6, 6, 6, 7, 7, 7, 7, 8, 8, 9, 10, 10,
        10, 11, 12, 12, 13, 13, 13, 13, 16, 16, 17, 17, 19, 19, 21, 21,
        25, 25, 27, 27, 27, 28, 29, 29, 30, 30, 31, 31, 31, 31, 32, 32,
        32, 32, 32, 32, 32, 32, 32, 32, 32, 33, 33, 33, 33, 33, 33, 33,
        33, 33, 33, 33, 33, 33, 33, 33, 33, 33])
    dst = np.array([0, 0, 1, 0, 1, 2, 0, 0, 0, 4, 5, 0, 1, 2, 3, 0, 2, 2, 0, 4,
        5, 0, 0, 3, 0, 1, 2, 3, 5, 6, 0, 1, 0, 1, 0, 1, 23, 24, 2, 23,
        24, 2, 23, 26, 1, 8, 0, 24, 25, 28, 2, 8, 14, 15, 18, 20, 22, 23,
        29, 30, 31, 8, 9, 13, 14, 15, 18, 19, 20, 22, 23, 26, 27, 28, 29, 30,
        31, 32])
    # Edges are directional in DGL; Make them bi-directional.
    u = np.concatenate([src, dst])
    v = np.concatenate([dst, src])
    # Construct a DGLGraph
    return dgl.graph((u, v))

class MyDataset(DGLDataset):
    def __init__(self,
                 url=None,
                 raw_dir=None,
                 save_dir=None,
                 force_reload=False,
                 verbose=False):
        super(MyDataset, self).__init__(name='dataset_name',
                                        url=url,
                                        raw_dir=raw_dir,
                                        save_dir=save_dir,
                                        force_reload=force_reload,
                                        verbose=verbose)

    def process(self):
        # 跳过一些处理的代码
        # === 跳过数据处理 ===

        # 构建图
        # g = dgl.graph(G)
        g = build_karate_club_graph()

        # train_mask = _sample_mask(idx_train, g.number_of_nodes())
        # val_mask = _sample_mask(idx_val, g.number_of_nodes())
        # test_mask = _sample_mask(idx_test, g.number_of_nodes())

        # # 划分掩码
        # g.ndata['train_mask'] = generate_mask_tensor(train_mask)
        # g.ndata['val_mask'] = generate_mask_tensor(val_mask)
        # g.ndata['test_mask'] = generate_mask_tensor(test_mask)

        # 节点的标签
        labels = torch.randint(0, 2, (g.number_of_edges(),))
        g.edata['labels'] = torch.tensor(labels)

        # 节点的特征
        g.ndata['features'] = torch.randn(g.number_of_nodes(), 10)
        self._num_labels = int(torch.max(labels).item() + 1)
        self._labels = labels
        self._g = g

    def __getitem__(self, idx):
        assert idx == 0, "这个数据集里只有一个图"
        return self._g

    def __len__(self):
        return 1

dataset = MyDataset()
g = dataset[0]
n_edges = g.number_of_edges()
train_eids = np.random.choice(np.arange(n_edges), (100,), replace=False)

定义邻居采样器

选择最简单的采样器:

sampler = dgl.dataloading.MultiLayerFullNeighborSampler(2)

在使用EdgeDataLoader()时有两种选择:1. 直接使用采样后的子图;2. 删除子图中出现在训练集中的边,防止模型使用这些边的信息。

直接使用采样后的子图

指定训练集中的边id即可,采样出来的结果依然是一些blocks:blocks[0]、blocks[1]… …,依次对应第0层到第1层的子图、第1层到第2层的子图… …。

train_eids = np.random.choice(np.arange(n_edges), (100,), replace=False)
dataloader = dgl.dataloading.EdgeDataLoader(
    g, train_eids, sampler,
    batch_size=10,
    shuffle=True,
    drop_last=False,
    num_workers=False)

删除子图中出现在训练集中的边

删除边的时候,除了正向的边,还有反向的边,所以使用reverse_eids指定需要删除的边,指定的方式,是一个边id构成的tensor。

dataloader = dgl.dataloading.EdgeDataLoader(
    g, train_eids, sampler,

    # 下面的两个参数专门用于在邻居采样时删除小批次的一些边和它们的反向边
    exclude='reverse_id',
    reverse_eids=torch.cat([
        torch.arange(n_edges // 2, n_edges), torch.arange(0, n_edges // 2)]),
    batch_size=10,
    shuffle=True,
    drop_last=False,
    num_workers=False)

获取节点表示

与节点分类的随机训练使用一样的模型:

class StochasticTwoLayerGCN(nn.Module):
    def __init__(self, in_features, hidden_features, out_features):
        super().__init__()
        self.conv1 = dglnn.GraphConv(in_features, hidden_features)
        self.conv2 = dglnn.GraphConv(hidden_features, out_features)

    def forward(self, blocks, x):
        x = F.relu(self.conv1(blocks[0], x))
        x = F.relu(self.conv2(blocks[1], x))
        return x

但是,如果采样时选择了删边的方式,那么就可能出现孤立的点,我们可以通过添加自环的方式来避免这个问题(g = dgl.add_self_loop(g)),也可以在模型中声明允许出现孤立点:“allow_zero_in_degree=True”。本文采用后者:

class StochasticTwoLayerGCN(nn.Module):
    def __init__(self, in_features, hidden_features, out_features):
        super().__init__()
        self.conv1 = dglnn.GraphConv(in_features, hidden_features, allow_zero_in_degree=True)
        self.conv2 = dglnn.GraphConv(hidden_features, out_features, allow_zero_in_degree=True)

    def forward(self, blocks, x):
        x = F.relu(self.conv1(blocks[0], x))
        x = F.relu(self.conv2(blocks[1], x))
        return x

预测边的得分

将边的两个端点的表示拼接起来,输入到一个全连接层得到边的得分。

class ScorePredictor(nn.Module):
    def __init__(self, num_classes, in_features):
        super().__init__()
        self.W = nn.Linear(2 * in_features, num_classes)

    def apply_edges(self, edges):
        data = torch.cat([edges.src['x'], edges.dst['x']], dim=-1)
        return {'score': self.W(data)}

    def forward(self, edge_subgraph, x):
        with edge_subgraph.local_scope():
            edge_subgraph.ndata['x'] = x
            edge_subgraph.apply_edges(self.apply_edges)
            return edge_subgraph.edata['score']

整体模型

就是分两步走,先通过两层GCN得到节点表示,然后经过边分数预测模型得到边的得分。

class Model(nn.Module):
    def __init__(self, in_features, hidden_features, out_features, num_classes):
        super().__init__()
        self.gcn = StochasticTwoLayerGCN(
            in_features, hidden_features, out_features)
        self.predictor = ScorePredictor(num_classes, out_features)

    def forward(self, edge_subgraph, blocks, x):
        x = self.gcn(blocks, x)
        return self.predictor(edge_subgraph, x)

开始训练

使用交叉熵作为损失函数。

def compute_loss(groundtruth, pred):
    return F.cross_entropy(pred, groundtruth)

in_features = 10
hidden_features = 100
out_features = 10
num_classes = 2
model = Model(in_features, hidden_features, out_features, num_classes)

opt = torch.optim.Adam(model.parameters())

for input_nodes, edge_subgraph, blocks in dataloader:
    input_features = blocks[0].srcdata['features']
    edge_labels = edge_subgraph.edata['labels']
    edge_predictions = model(edge_subgraph, blocks, input_features)
    loss = compute_loss(edge_labels, edge_predictions)
    opt.zero_grad()
    loss.backward()
    print('loss: ', loss.item())
    opt.step()

随机训练之异构图上的边分类

加载数据

还是使用“跟着官方文档学DGL框架第八天”中人工构建的异构图数据集。
为了简便,我们只针对“click”类型的边做分类任务,随机给这些边打上标签。

n_users = 1000
n_items = 500
n_follows = 3000
n_clicks = 5000
n_dislikes = 500
n_hetero_features = 10
n_user_classes = 5
n_max_clicks = 10

follow_src = np.random.randint(0, n_users, n_follows)
follow_dst = np.random.randint(0, n_users, n_follows)
click_src = np.random.randint(0, n_users, n_clicks)
click_dst = np.random.randint(0, n_items, n_clicks)
dislike_src = np.random.randint(0, n_users, n_dislikes)
dislike_dst = np.random.randint(0, n_items, n_dislikes)

hetero_graph = dgl.heterograph({
    ('user', 'follow', 'user'): (follow_src, follow_dst),
    ('user', 'followed-by', 'user'): (follow_dst, follow_src),
    ('user', 'click', 'item'): (click_src, click_dst),
    ('item', 'clicked-by', 'user'): (click_dst, click_src),
    ('user', 'dislike', 'item'): (dislike_src, dislike_dst),
    ('item', 'disliked-by', 'user'): (dislike_dst, dislike_src)})

hetero_graph.nodes['user'].data['feat'] = torch.randn(n_users, n_hetero_features)
hetero_graph.nodes['item'].data['feat'] = torch.randn(n_items, n_hetero_features)
hetero_graph.edges['click'].data['labels'] = torch.randint(0, 2, (n_clicks,))

g = hetero_graph

train_eid_dict = {'click': np.random.choice(np.arange(n_clicks), (1000, ), replace=False)}

定义邻居采样器

这里只给了删除边的方式。由于是异构图,正向边和反向边具有不同的类型,此时“exclude=‘reverse_types’”,而且在“reverse_etypes”里需要指出正向边类型与反向边类型的映射字典。

sampler = dgl.dataloading.MultiLayerFullNeighborSampler(2)
dataloader = dgl.dataloading.EdgeDataLoader(
    g, train_eid_dict, sampler,

    # 下面的两个参数专门用于在邻居采样时删除小批次的一些边和它们的反向边
    exclude='reverse_types',
    reverse_etypes={'follow': 'followed-by', 'followed-by': 'follow',
                    'click': 'clicked-by', 'clicked-by': 'click',
                    'dislike': 'disliked-by', 'disliked-by': 'dislike'},

    batch_size=10,
    shuffle=True,
    drop_last=False,
    num_workers=False)

获取节点表示

与节点分类的随机训练使用一样的模型:

class StochasticTwoLayerRGCN(nn.Module):
    def __init__(self, in_feat, hidden_feat, out_feat, rel_names):
        super().__init__()
        self.conv1 = dglnn.HeteroGraphConv({
                rel : dglnn.GraphConv(in_feat, hidden_feat, norm='right', allow_zero_in_degree=True)
                for rel in rel_names
            })
        self.conv2 = dglnn.HeteroGraphConv({
                rel : dglnn.GraphConv(hidden_feat, out_feat, norm='right', allow_zero_in_degree=True)
                for rel in rel_names
            })

    def forward(self, blocks, x):
        x = self.conv1(blocks[0], x)
        x = self.conv2(blocks[1], x)
        return x

预测边的得分

与同构图时的区别在于,获取边得分时需要分不同的边类型。

class ScorePredictor(nn.Module):
    def __init__(self, num_classes, in_features):
        super().__init__()
        self.W = nn.Linear(2 * in_features, num_classes)

    def apply_edges(self, edges):
        data = torch.cat([edges.src['x'], edges.dst['x']], dim=-1)
        return {'score': self.W(data)}

    def forward(self, edge_subgraph, x):
        with edge_subgraph.local_scope():
            edge_subgraph.ndata['x'] = x
            for etype in edge_subgraph.canonical_etypes:
                edge_subgraph.apply_edges(self.apply_edges, etype=etype)
            return edge_subgraph.edata['score']

整体模型

与同构图时没有太大区别,只是换成了异构图上的卷积模型。

class Model(nn.Module):
    def __init__(self, in_features, hidden_features, out_features, num_classes,
                 etypes):
        super().__init__()
        self.rgcn = StochasticTwoLayerRGCN(
            in_features, hidden_features, out_features, etypes)
        self.pred = ScorePredictor(num_classes, out_features)

    def forward(self, edge_subgraph, blocks, x):
        x = self.rgcn(blocks, x)
        return self.pred(edge_subgraph, x)

开始训练

依然采用交叉熵损失函数,只是注意模型的输出是字典形式的,键为边的类型,值为边的得分。

def compute_loss(groundtruth, pred):
    return F.cross_entropy(pred, groundtruth)

in_features = n_hetero_features
hidden_features = 100
out_features = 10
num_classes = 2
etypes = g.etypes

model = Model(in_features, hidden_features, out_features, num_classes, etypes)
opt = torch.optim.Adam(model.parameters())
for input_nodes, edge_subgraph, blocks in dataloader:
    input_features = blocks[0].srcdata['feat']
    edge_labels = edge_subgraph.edata['labels']
    edge_predictions = model(edge_subgraph, blocks, input_features)
    loss = compute_loss(edge_labels[('user', 'click', 'item')], edge_predictions[('user', 'click', 'item')])
    opt.zero_grad()
    loss.backward()
    print('loss: ', loss.item())
    opt.step()

完整代码

随机训练之同构图上的边分类

import dgl
import dgl.nn as dglnn
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from dgl.data.utils import generate_mask_tensor
from dgl.data import DGLDataset
import torch

def build_karate_club_graph():
    # All 78 edges are stored in two numpy arrays. One for source endpoints
    # while the other for destination endpoints.
    src = np.array([1, 2, 2, 3, 3, 3, 4, 5, 6, 6, 6, 7, 7, 7, 7, 8, 8, 9, 10, 10,
        10, 11, 12, 12, 13, 13, 13, 13, 16, 16, 17, 17, 19, 19, 21, 21,
        25, 25, 27, 27, 27, 28, 29, 29, 30, 30, 31, 31, 31, 31, 32, 32,
        32, 32, 32, 32, 32, 32, 32, 32, 32, 33, 33, 33, 33, 33, 33, 33,
        33, 33, 33, 33, 33, 33, 33, 33, 33, 33])
    dst = np.array([0, 0, 1, 0, 1, 2, 0, 0, 0, 4, 5, 0, 1, 2, 3, 0, 2, 2, 0, 4,
        5, 0, 0, 3, 0, 1, 2, 3, 5, 6, 0, 1, 0, 1, 0, 1, 23, 24, 2, 23,
        24, 2, 23, 26, 1, 8, 0, 24, 25, 28, 2, 8, 14, 15, 18, 20, 22, 23,
        29, 30, 31, 8, 9, 13, 14, 15, 18, 19, 20, 22, 23, 26, 27, 28, 29, 30,
        31, 32])
    # Edges are directional in DGL; Make them bi-directional.
    u = np.concatenate([src, dst])
    v = np.concatenate([dst, src])
    # Construct a DGLGraph
    return dgl.graph((u, v))

# def _sample_mask(idx, l):
#     """Create mask."""
#     mask = np.zeros(l)
#     mask[idx] = 1
#     return mask

class MyDataset(DGLDataset):
    def __init__(self,
                 url=None,
                 raw_dir=None,
                 save_dir=None,
                 force_reload=False,
                 verbose=False):
        super(MyDataset, self).__init__(name='dataset_name',
                                        url=url,
                                        raw_dir=raw_dir,
                                        save_dir=save_dir,
                                        force_reload=force_reload,
                                        verbose=verbose)

    def process(self):
        # 跳过一些处理的代码
        # === 跳过数据处理 ===

        # 构建图
        # g = dgl.graph(G)
        g = build_karate_club_graph()

        # train_mask = _sample_mask(idx_train, g.number_of_nodes())
        # val_mask = _sample_mask(idx_val, g.number_of_nodes())
        # test_mask = _sample_mask(idx_test, g.number_of_nodes())

        # # 划分掩码
        # g.ndata['train_mask'] = generate_mask_tensor(train_mask)
        # g.ndata['val_mask'] = generate_mask_tensor(val_mask)
        # g.ndata['test_mask'] = generate_mask_tensor(test_mask)

        # 节点的标签
        labels = torch.randint(0, 2, (g.number_of_edges(),))
        g.edata['labels'] = torch.tensor(labels)

        # 节点的特征
        g.ndata['features'] = torch.randn(g.number_of_nodes(), 10)
        self._num_labels = int(torch.max(labels).item() + 1)
        self._labels = labels
        self._g = g

    def __getitem__(self, idx):
        assert idx == 0, "这个数据集里只有一个图"
        return self._g

    def __len__(self):
        return 1

dataset = MyDataset()
g = dataset[0]
n_edges = g.number_of_edges()
train_eids = np.random.choice(np.arange(n_edges), (100,), replace=False)
sampler = dgl.dataloading.MultiLayerFullNeighborSampler(2)

# dataloader = dgl.dataloading.EdgeDataLoader(
#     g, train_eids, sampler,
#     batch_size=10,
#     shuffle=True,
#     drop_last=False,
#     num_workers=False)


dataloader = dgl.dataloading.EdgeDataLoader(
    g, train_eids, sampler,

    # 下面的两个参数专门用于在邻居采样时删除小批次的一些边和它们的反向边
    exclude='reverse_id',
    reverse_eids=torch.cat([
        torch.arange(n_edges // 2, n_edges), torch.arange(0, n_edges // 2)]),
    batch_size=10,
    shuffle=True,
    drop_last=False,
    num_workers=False)

class StochasticTwoLayerGCN(nn.Module):
    def __init__(self, in_features, hidden_features, out_features):
        super().__init__()
        self.conv1 = dglnn.GraphConv(in_features, hidden_features, allow_zero_in_degree=True)
        self.conv2 = dglnn.GraphConv(hidden_features, out_features, allow_zero_in_degree=True)

    def forward(self, blocks, x):
        x = F.relu(self.conv1(blocks[0], x))
        x = F.relu(self.conv2(blocks[1], x))
        return x

class ScorePredictor(nn.Module):
    def __init__(self, num_classes, in_features):
        super().__init__()
        self.W = nn.Linear(2 * in_features, num_classes)

    def apply_edges(self, edges):
        data = torch.cat([edges.src['x'], edges.dst['x']], dim=-1)
        return {'score': self.W(data)}

    def forward(self, edge_subgraph, x):
        with edge_subgraph.local_scope():
            edge_subgraph.ndata['x'] = x
            edge_subgraph.apply_edges(self.apply_edges)
            return edge_subgraph.edata['score']


class Model(nn.Module):
    def __init__(self, in_features, hidden_features, out_features, num_classes):
        super().__init__()
        self.gcn = StochasticTwoLayerGCN(
            in_features, hidden_features, out_features)
        self.predictor = ScorePredictor(num_classes, out_features)

    def forward(self, edge_subgraph, blocks, x):
        x = self.gcn(blocks, x)
        return self.predictor(edge_subgraph, x)

def compute_loss(groundtruth, pred):
    return F.cross_entropy(pred, groundtruth)

in_features = 10
hidden_features = 100
out_features = 10
num_classes = 2
model = Model(in_features, hidden_features, out_features, num_classes)
# model = model.cuda()
# opt = torch.optim.Adam(model.parameters())

# for input_nodes, edge_subgraph, blocks in dataloader:
#     blocks = [b.to(torch.device('cuda')) for b in blocks]
#     edge_subgraph = edge_subgraph.to(torch.device('cuda'))
#     input_features = blocks[0].srcdata['features']
#     edge_labels = edge_subgraph.edata['labels']
#     edge_predictions = model(edge_subgraph, blocks, input_features)
#     loss = compute_loss(edge_labels, edge_predictions)
#     opt.zero_grad()
#     loss.backward()
#     opt.step()

opt = torch.optim.Adam(model.parameters())

for input_nodes, edge_subgraph, blocks in dataloader:
    input_features = blocks[0].srcdata['features']
    edge_labels = edge_subgraph.edata['labels']
    edge_predictions = model(edge_subgraph, blocks, input_features)
    loss = compute_loss(edge_labels, edge_predictions)
    opt.zero_grad()
    loss.backward()
    print('loss: ', loss.item())
    opt.step()

随机训练之异构图上的边分类

import dgl
import dgl.nn as dglnn
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

n_users = 1000
n_items = 500
n_follows = 3000
n_clicks = 5000
n_dislikes = 500
n_hetero_features = 10
n_user_classes = 5
n_max_clicks = 10

follow_src = np.random.randint(0, n_users, n_follows)
follow_dst = np.random.randint(0, n_users, n_follows)
click_src = np.random.randint(0, n_users, n_clicks)
click_dst = np.random.randint(0, n_items, n_clicks)
dislike_src = np.random.randint(0, n_users, n_dislikes)
dislike_dst = np.random.randint(0, n_items, n_dislikes)

hetero_graph = dgl.heterograph({
    ('user', 'follow', 'user'): (follow_src, follow_dst),
    ('user', 'followed-by', 'user'): (follow_dst, follow_src),
    ('user', 'click', 'item'): (click_src, click_dst),
    ('item', 'clicked-by', 'user'): (click_dst, click_src),
    ('user', 'dislike', 'item'): (dislike_src, dislike_dst),
    ('item', 'disliked-by', 'user'): (dislike_dst, dislike_src)})

hetero_graph.nodes['user'].data['feat'] = torch.randn(n_users, n_hetero_features)
hetero_graph.nodes['item'].data['feat'] = torch.randn(n_items, n_hetero_features)
hetero_graph.edges['click'].data['labels'] = torch.randint(0, 2, (n_clicks,))

g = hetero_graph

train_eid_dict = {'click': np.random.choice(np.arange(n_clicks), (1000, ), replace=False)}

class StochasticTwoLayerRGCN(nn.Module):
    def __init__(self, in_feat, hidden_feat, out_feat, rel_names):
        super().__init__()
        self.conv1 = dglnn.HeteroGraphConv({
                rel : dglnn.GraphConv(in_feat, hidden_feat, norm='right', allow_zero_in_degree=True)
                for rel in rel_names
            })
        self.conv2 = dglnn.HeteroGraphConv({
                rel : dglnn.GraphConv(hidden_feat, out_feat, norm='right', allow_zero_in_degree=True)
                for rel in rel_names
            })

    def forward(self, blocks, x):
        x = self.conv1(blocks[0], x)
        x = self.conv2(blocks[1], x)
        return x

class ScorePredictor(nn.Module):
    def __init__(self, num_classes, in_features):
        super().__init__()
        self.W = nn.Linear(2 * in_features, num_classes)

    def apply_edges(self, edges):
        data = torch.cat([edges.src['x'], edges.dst['x']], dim=-1)
        return {'score': self.W(data)}

    def forward(self, edge_subgraph, x):
        with edge_subgraph.local_scope():
            edge_subgraph.ndata['x'] = x
            for etype in edge_subgraph.canonical_etypes:
                edge_subgraph.apply_edges(self.apply_edges, etype=etype)
            return edge_subgraph.edata['score']

class Model(nn.Module):
    def __init__(self, in_features, hidden_features, out_features, num_classes,
                 etypes):
        super().__init__()
        self.rgcn = StochasticTwoLayerRGCN(
            in_features, hidden_features, out_features, etypes)
        self.pred = ScorePredictor(num_classes, out_features)

    def forward(self, edge_subgraph, blocks, x):
        x = self.rgcn(blocks, x)
        return self.pred(edge_subgraph, x)

sampler = dgl.dataloading.MultiLayerFullNeighborSampler(2)
# dataloader = dgl.dataloading.EdgeDataLoader(
#     g, train_eid_dict, sampler,
#     batch_size=10,
#     shuffle=True,
#     drop_last=False,
#     num_workers=False)

dataloader = dgl.dataloading.EdgeDataLoader(
    g, train_eid_dict, sampler,

    # 下面的两个参数专门用于在邻居采样时删除小批次的一些边和它们的反向边
    exclude='reverse_types',
    reverse_etypes={'follow': 'followed-by', 'followed-by': 'follow',
                    'click': 'clicked-by', 'clicked-by': 'click',
                    'dislike': 'disliked-by', 'disliked-by': 'dislike'},

    batch_size=10,
    shuffle=True,
    drop_last=False,
    num_workers=False)
def compute_loss(groundtruth, pred):
    return F.cross_entropy(pred, groundtruth)

in_features = n_hetero_features
hidden_features = 100
out_features = 10
num_classes = 2
etypes = g.etypes

model = Model(in_features, hidden_features, out_features, num_classes, etypes)
# model = model.cuda()
# opt = torch.optim.Adam(model.parameters())

# for input_nodes, edge_subgraph, blocks in dataloader:
#     blocks = [b.to(torch.device('cuda')) for b in blocks]
#     edge_subgraph = edge_subgraph.to(torch.device('cuda'))
#     input_features = blocks[0].srcdata['features']
#     edge_labels = edge_subgraph.edata['labels']
#     edge_predictions = model(edge_subgraph, blocks, input_features)
#     loss = compute_loss(edge_labels, edge_predictions)
#     opt.zero_grad()
#     loss.backward()
#     opt.step()

opt = torch.optim.Adam(model.parameters())
for input_nodes, edge_subgraph, blocks in dataloader:
    input_features = blocks[0].srcdata['feat']
    edge_labels = edge_subgraph.edata['labels']
    edge_predictions = model(edge_subgraph, blocks, input_features)
    loss = compute_loss(edge_labels[('user', 'click', 'item')], edge_predictions[('user', 'click', 'item')])
    opt.zero_grad()
    loss.backward()
    print('loss: ', loss.item())
    opt.step()

你可能感兴趣的:(DGL,深度学习,神经网络,图神经网络,DGL,边分类)