跟着官方文档学DGL框架第十四天——大型图上的随机训练之链接预测

参考资料

1. https://docs.dgl.ai/en/latest/guide/minibatch-link.html#guide-minibatch-link-classification-sampler
2. http://www.voidcn.com/article/p-fsdktdik-bry.html

概述

关于链接预测的概念和优化方法在“跟着官方文档学DGL框架第十天”上已经提到过。我们的目标还是得到节点表示，所以在随机训练时与节点分类和边分类的随机训练差不多，只是多了负采样过程。值得庆幸的是，DGL在随机训练时的负采样，只需要指定dgl.dataloading.EdgeDataLoader()中的negative_sampler为你需要的负采样函数。

随机训练之同构图上的链接预测

加载数据

依然使用“跟着官方文档学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_seeds = np.random.choice(np.arange(n_edges), (20,), replace=False)

定义邻居采样器

依然选择最简单的采样器：

sampler = dgl.dataloading.MultiLayerFullNeighborSampler(2)

负采样时有两种方式，一是使用DGL自带的随机采样，二是自定义负采样函数。

随机采样

随机采样只需要指定“negative_sampler=dgl.dataloading.negative_sampler.Uniform(5)”，其中“5”表示负样本个数。

“drop_last”和“pin_memory”参数来自torch.data.DataLoader。“drop_last”表示是否去除最后一个不完整的batch；“pin_memory”表示是否使用锁页内存，建议在使用GPU时设置为True。

dataloader会返回“输入节点”、“正样本图”、“负采样图”和“子图块”四个结果。

dataloader = dgl.dataloading.EdgeDataLoader(
    g, train_seeds, sampler,
    negative_sampler=dgl.dataloading.negative_sampler.Uniform(5),
    batch_size=4,
    shuffle=True,
    drop_last=False,
    pin_memory=True,
    num_workers=False)

自定义负采样函数

负采样函数初始化的参数为：1.原始图“g”；2.一个正样本对应的负样本个数“k”。函数的输入为原始图“g”和小批量的边id“eids”，返回的结果是负样本的源节点数组和目标节点数组。

下面是按原始图中节点度的0.75次幂为采样率的例子。

class NegativeSampler(object):
    def __init__(self, g, k):
        # caches the probability distribution
        self.weights = g.in_degrees().float() ** 0.75
        self.k = k

    def __call__(self, g, eids):
        src, _ = g.find_edges(eids)
        src = src.repeat_interleave(self.k)
        dst = self.weights.multinomial(len(src), replacement=True)
        return src, dst

dataloader = dgl.dataloading.EdgeDataLoader(
    g, train_seeds, sampler,
    negative_sampler=NegativeSampler(g, 5),
    batch_size=4,
    shuffle=True,
    drop_last=False,
    pin_memory=True,
    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

预测边的得分

这里使用边的两个端点的内积作为分数：

class ScorePredictor(nn.Module):
    def forward(self, edge_subgraph, x):
        with edge_subgraph.local_scope():
            edge_subgraph.ndata['x'] = x
            edge_subgraph.apply_edges(dgl.function.u_dot_v('x', 'x', 'score'))
            return edge_subgraph.edata['score']

整体模型

首先获得节点的表示，然后分别计算正样本图和负采样图上的边得分。

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

    def forward(self, positive_graph, negative_graph, blocks, x):
        x = self.gcn(blocks, x)
        pos_score = self.predictor(positive_graph, x)
        neg_score = self.predictor(negative_graph, x)
        return pos_score, neg_score

开始训练

使用hinge loss作为损失函数。

def compute_loss(pos_score, neg_score):
    # an example hinge loss
    n = pos_score.shape[0]
    return (neg_score.view(n, -1) - pos_score.view(n, -1) + 1).clamp(min=0).mean()

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

for input_nodes, positive_graph, negative_graph, blocks in dataloader:
    input_features = blocks[0].srcdata['features']
    pos_score, neg_score = model(positive_graph, negative_graph, blocks, input_features)
    loss = compute_loss(pos_score, neg_score)
    opt.zero_grad()
    loss.backward()
    print('loss: ', loss.item())
    opt.step()

随机训练之异构图上的链接预测

加载数据

还是使用“跟着官方文档学DGL框架第八天”中人工构建的异构图数据集。训练集选择了所有类型的所有边，以字典的形式给出。

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)

g = hetero_graph

train_eid_dict = {
    etype: g.edges(etype=etype, form='eid')
    for etype in g.etypes}

定义邻居采样器

依然选择最简单的采样器：

sampler = dgl.dataloading.MultiLayerFullNeighborSampler(2)

负采样时有两种方式，一是使用DGL自带的随机采样，二是自定义负采样函数。

随机采样

与同构图时无异，dgl.dataloading.negative_sampler.Uniform()同样支持异构图。

sampler = dgl.dataloading.MultiLayerFullNeighborSampler(2)

dataloader = dgl.dataloading.EdgeDataLoader(
    g, train_eid_dict, sampler,
    negative_sampler=dgl.dataloading.negative_sampler.Uniform(5),
    batch_size=10,
    shuffle=True,
    drop_last=False,
    num_workers=False)

自定义负采样函数

这个还没有调通，之后再来调。[flag]

class NegativeSampler(object):
    def __init__(self, g, k):
        # 缓存概率分布
        self.weights = {
            etype: g.in_degrees(etype=etype).float() ** 0.75
            for _, etype, _ in g.canonical_etypes
        }
        self.k = k

    def __call__(self, g, eids_dict):
        result_dict = {}
        for etype, eids in eids_dict.items():
            src, _ = g.find_edges(eids, etype=etype)
            src = src.repeat_interleave(self.k)
            dst = self.weights[etype].multinomial(len(src), replacement=True)
            result_dict[etype] = (src, dst)
        return result_dict

获取节点表示

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

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')
                for rel in rel_names
            })
        self.conv2 = dglnn.HeteroGraphConv({
                rel : dglnn.GraphConv(hidden_feat, out_feat, norm='right')
                for rel in rel_names
            })

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

预测边的得分

这里使用边的两个端点的内积作为分数，与同构图的区别在于，需要分边类型执行apply_edges()：

class ScorePredictor(nn.Module):
    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(
                    dgl.function.u_dot_v('x', 'x', 'score'), etype=etype)
            return edge_subgraph.edata['score']

整体模型

首先获得节点的表示，然后分别计算正样本图和负采样图上的边得分，注意返回的结果是字典形式，键是边类型，值是分数。

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

    def forward(self, positive_graph, negative_graph, blocks, x):
        x = self.rgcn(blocks, x)
        pos_score = self.pred(positive_graph, x)
        neg_score = self.pred(negative_graph, x)
        return pos_score, neg_score

开始训练

由于返回的分数结果是字典形式，所以需要自定义一个损失函数，这里对每种边类型分别使用hinge loss，再求和作为最终损失。

def compute_loss(pos_score, neg_score):
    loss = 0
    # an example hinge loss
    for etype, p_score in pos_score.items():
        if len(p_score) != 0:
            n = p_score.shape[0]
            loss += (neg_score[etype].view(n, -1) - p_score.view(n, -1) + 1).clamp(min=0).mean()
        
    return loss

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

model = Model(in_features, hidden_features, out_features, etypes)

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

for input_nodes, positive_graph, negative_graph, blocks in dataloader:
    print('negative graph: ', negative_graph)
    input_features = blocks[0].srcdata['feat']
    pos_score, neg_score = model(positive_graph, negative_graph, blocks, input_features)
    loss = compute_loss(pos_score, neg_score)
    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_seeds = np.random.choice(np.arange(n_edges), (20,), replace=False)

sampler = dgl.dataloading.MultiLayerFullNeighborSampler(2)

dataloader = dgl.dataloading.EdgeDataLoader(
    g, train_seeds, sampler,
    negative_sampler=dgl.dataloading.negative_sampler.Uniform(5),
    batch_size=4,
    shuffle=True,
    drop_last=False,
    pin_memory=True,
    num_workers=False)


# class NegativeSampler(object):
#     def __init__(self, g, k):
#         # caches the probability distribution
#         self.weights = g.in_degrees().float() ** 0.75
#         self.k = k

#     def __call__(self, g, eids):
#         src, _ = g.find_edges(eids)
#         src = src.repeat_interleave(self.k)
#         dst = self.weights.multinomial(len(src), replacement=True)
#         return src, dst

# dataloader = dgl.dataloading.EdgeDataLoader(
#     g, train_seeds, sampler,
#     negative_sampler=NegativeSampler(g, 5),
#     batch_size=4,
#     shuffle=True,
#     drop_last=False,
#     pin_memory=True,
#     num_workers=False)


class StochasticTwoLayerGCN(nn.Module):
    def __init__(self, in_features, hidden_features, out_features):
        super().__init__()
        self.conv1 = dgl.nn.GraphConv(in_features, hidden_features)
        self.conv2 = dgl.nn.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


class ScorePredictor(nn.Module):
    def forward(self, edge_subgraph, x):
        with edge_subgraph.local_scope():
            edge_subgraph.ndata['x'] = x
            edge_subgraph.apply_edges(dgl.function.u_dot_v('x', 'x', 'score'))
            return edge_subgraph.edata['score']


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

    def forward(self, positive_graph, negative_graph, blocks, x):
        x = self.gcn(blocks, x)
        pos_score = self.predictor(positive_graph, x)
        neg_score = self.predictor(negative_graph, x)
        return pos_score, neg_score

def compute_loss(pos_score, neg_score):
    # an example hinge loss
    n = pos_score.shape[0]
    return (neg_score.view(n, -1) - pos_score.view(n, -1) + 1).clamp(min=0).mean()

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


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

# for input_nodes, positive_graph, negative_graph, blocks in dataloader:
#     blocks = [b.to(torch.device('cuda')) for b in blocks]
#     positive_graph = positive_graph.to(torch.device('cuda'))
#     negative_graph = negative_graph.to(torch.device('cuda'))
#     input_features = blocks[0].srcdata['features']
#     pos_score, neg_score = model(positive_graph, negative_graph, blocks, input_features)
#     loss = compute_loss(pos_score, neg_score)
#     opt.zero_grad()
#     loss.backward()
#     opt.step()

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

for input_nodes, positive_graph, negative_graph, blocks in dataloader:
    input_features = blocks[0].srcdata['features']
    pos_score, neg_score = model(positive_graph, negative_graph, blocks, input_features)
    loss = compute_loss(pos_score, neg_score)
    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

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)

g = hetero_graph

train_eid_dict = {
    etype: g.edges(etype=etype, form='eid')
    for etype in g.etypes}

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')
                for rel in rel_names
            })
        self.conv2 = dglnn.HeteroGraphConv({
                rel : dglnn.GraphConv(hidden_feat, out_feat, norm='right')
                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 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(
                    dgl.function.u_dot_v('x', 'x', 'score'), etype=etype)
            return edge_subgraph.edata['score']

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

    def forward(self, positive_graph, negative_graph, blocks, x):
        x = self.rgcn(blocks, x)
        pos_score = self.pred(positive_graph, x)
        neg_score = self.pred(negative_graph, x)
        return pos_score, neg_score


sampler = dgl.dataloading.MultiLayerFullNeighborSampler(2)

dataloader = dgl.dataloading.EdgeDataLoader(
    g, train_eid_dict, sampler,
    negative_sampler=dgl.dataloading.negative_sampler.Uniform(5),
    batch_size=10,
    shuffle=True,
    drop_last=False,
    num_workers=False)

# class NegativeSampler(object):
#     def __init__(self, g, k):
#         # 缓存概率分布
#         self.weights = {
#             etype: g.in_degrees(etype=etype).float() ** 0.75
#             for _, etype, _ in g.canonical_etypes
#         }
#         self.k = k

#     def __call__(self, g, eids_dict):
#         result_dict = {}
#         for etype, eids in eids_dict.items():
#             print(etype)
#             src, _ = g.find_edges(eids, etype=etype)
#             src = src.repeat_interleave(self.k)
#             dst = self.weights[etype].multinomial(len(src), replacement=True)
#             result_dict[etype] = (src, dst)
#             print('len_dict: ', result_dict[etype])
#         return result_dict


# dataloader = dgl.dataloading.EdgeDataLoader(
#     g, train_eid_dict, sampler,
#     negative_sampler=NegativeSampler(g, 5),
#     batch_size=1000,
#     shuffle=True,
#     drop_last=False,
#     num_workers=False)

def compute_loss(pos_score, neg_score):
    loss = 0
    # an example hinge loss
    for etype, p_score in pos_score.items():
        if len(p_score) != 0:
            n = p_score.shape[0]
            loss += (neg_score[etype].view(n, -1) - p_score.view(n, -1) + 1).clamp(min=0).mean()
        
    return loss

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

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

# for input_nodes, positive_graph, negative_graph, blocks in dataloader:
#     blocks = [b.to(torch.device('cuda')) for b in blocks]
#     positive_graph = positive_graph.to(torch.device('cuda'))
#     negative_graph = negative_graph.to(torch.device('cuda'))
#     input_features = blocks[0].srcdata['features']
#     pos_score, neg_score = model(positive_graph, negative_graph, blocks, input_features)
#     loss = compute_loss(pos_score, neg_score)
#     opt.zero_grad()
#     loss.backward()
#     print('loss: ', loss.item())
#     opt.step()

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

for input_nodes, positive_graph, negative_graph, blocks in dataloader:
    print('negative graph: ', negative_graph)
    input_features = blocks[0].srcdata['feat']
    pos_score, neg_score = model(positive_graph, negative_graph, blocks, input_features)
    loss = compute_loss(pos_score, neg_score)
    opt.zero_grad()
    loss.backward()
    print('loss: ', loss.item())
    opt.step()