Source code for dgl.nn.pytorch.network_emb
"""Network Embedding NN Modules"""
# pylint: disable= invalid-name
import random
import torch
import torch.nn.functional as F
from torch import nn
from torch.nn import init
from tqdm.auto import trange
from ...base import NID
from ...convert import to_heterogeneous, to_homogeneous
from ...random import choice
from ...sampling import random_walk
__all__ = ["DeepWalk", "MetaPath2Vec"]
[docs]class DeepWalk(nn.Module):
"""DeepWalk module from `DeepWalk: Online Learning of Social Representations
<https://arxiv.org/abs/1403.6652>`__
For a graph, it learns the node representations from scratch by maximizing the similarity of
node pairs that are nearby (positive node pairs) and minimizing the similarity of other
random node pairs (negative node pairs).
Parameters
----------
g : DGLGraph
Graph for learning node embeddings
emb_dim : int, optional
Size of each embedding vector. Default: 128
walk_length : int, optional
Number of nodes in a random walk sequence. Default: 40
window_size : int, optional
In a random walk :attr:`w`, a node :attr:`w[j]` is considered close to a node
:attr:`w[i]` if :attr:`i - window_size <= j <= i + window_size`. Default: 5
neg_weight : float, optional
Weight of the loss term for negative samples in the total loss. Default: 1.0
negative_size : int, optional
Number of negative samples to use for each positive sample. Default: 5
fast_neg : bool, optional
If True, it samples negative node pairs within a batch of random walks. Default: True
sparse : bool, optional
If True, gradients with respect to the learnable weights will be sparse.
Default: True
Attributes
----------
node_embed : nn.Embedding
Embedding table of the nodes
Examples
--------
>>> import torch
>>> from dgl.data import CoraGraphDataset
>>> from dgl.nn import DeepWalk
>>> from torch.optim import SparseAdam
>>> from torch.utils.data import DataLoader
>>> from sklearn.linear_model import LogisticRegression
>>> dataset = CoraGraphDataset()
>>> g = dataset[0]
>>> model = DeepWalk(g)
>>> dataloader = DataLoader(torch.arange(g.num_nodes()), batch_size=128,
... shuffle=True, collate_fn=model.sample)
>>> optimizer = SparseAdam(model.parameters(), lr=0.01)
>>> num_epochs = 5
>>> for epoch in range(num_epochs):
... for batch_walk in dataloader:
... loss = model(batch_walk)
... optimizer.zero_grad()
... loss.backward()
... optimizer.step()
>>> train_mask = g.ndata['train_mask']
>>> test_mask = g.ndata['test_mask']
>>> X = model.node_embed.weight.detach()
>>> y = g.ndata['label']
>>> clf = LogisticRegression().fit(X[train_mask].numpy(), y[train_mask].numpy())
>>> clf.score(X[test_mask].numpy(), y[test_mask].numpy())
"""
def __init__(
self,
g,
emb_dim=128,
walk_length=40,
window_size=5,
neg_weight=1,
negative_size=5,
fast_neg=True,
sparse=True,
):
super().__init__()
assert (
walk_length >= window_size + 1
), f"Expect walk_length >= window_size + 1, got {walk_length} and {window_size + 1}"
self.g = g
self.emb_dim = emb_dim
self.window_size = window_size
self.walk_length = walk_length
self.neg_weight = neg_weight
self.negative_size = negative_size
self.fast_neg = fast_neg
num_nodes = g.num_nodes()
# center node embedding
self.node_embed = nn.Embedding(num_nodes, emb_dim, sparse=sparse)
self.context_embed = nn.Embedding(num_nodes, emb_dim, sparse=sparse)
self.reset_parameters()
if not fast_neg:
neg_prob = g.out_degrees().pow(0.75)
# categorical distribution for true negative sampling
self.neg_prob = neg_prob / neg_prob.sum()
# Get list index pairs for positive samples.
# Given i, positive index pairs are (i - window_size, i), ... ,
# (i - 1, i), (i + 1, i), ..., (i + window_size, i)
idx_list_src = []
idx_list_dst = []
for i in range(walk_length):
for j in range(max(0, i - window_size), i):
idx_list_src.append(j)
idx_list_dst.append(i)
for j in range(i + 1, min(walk_length, i + 1 + window_size)):
idx_list_src.append(j)
idx_list_dst.append(i)
self.idx_list_src = torch.LongTensor(idx_list_src)
self.idx_list_dst = torch.LongTensor(idx_list_dst)
[docs] def reset_parameters(self):
"""Reinitialize learnable parameters"""
init_range = 1.0 / self.emb_dim
init.uniform_(self.node_embed.weight.data, -init_range, init_range)
init.constant_(self.context_embed.weight.data, 0)
def sample(self, indices):
"""Sample random walks
Parameters
----------
indices : torch.Tensor
Nodes from which we perform random walk
Returns
-------
torch.Tensor
Random walks in the form of node ID sequences. The Tensor
is of shape :attr:`(len(indices), walk_length)`.
"""
return random_walk(self.g, indices, length=self.walk_length - 1)[0]
[docs] def forward(self, batch_walk):
"""Compute the loss for the batch of random walks
Parameters
----------
batch_walk : torch.Tensor
Random walks in the form of node ID sequences. The Tensor
is of shape :attr:`(batch_size, walk_length)`.
Returns
-------
torch.Tensor
Loss value
"""
batch_size = len(batch_walk)
device = batch_walk.device
batch_node_embed = self.node_embed(batch_walk).view(-1, self.emb_dim)
batch_context_embed = self.context_embed(batch_walk).view(
-1, self.emb_dim
)
batch_idx_list_offset = torch.arange(batch_size) * self.walk_length
batch_idx_list_offset = batch_idx_list_offset.unsqueeze(1)
idx_list_src = batch_idx_list_offset + self.idx_list_src.unsqueeze(0)
idx_list_dst = batch_idx_list_offset + self.idx_list_dst.unsqueeze(0)
idx_list_src = idx_list_src.view(-1).to(device)
idx_list_dst = idx_list_dst.view(-1).to(device)
pos_src_emb = batch_node_embed[idx_list_src]
pos_dst_emb = batch_context_embed[idx_list_dst]
neg_idx_list_src = idx_list_dst.unsqueeze(1) + torch.zeros(
self.negative_size
).unsqueeze(0).to(device)
neg_idx_list_src = neg_idx_list_src.view(-1)
neg_src_emb = batch_node_embed[neg_idx_list_src.long()]
if self.fast_neg:
neg_idx_list_dst = list(range(batch_size * self.walk_length)) * (
self.negative_size * self.window_size * 2
)
random.shuffle(neg_idx_list_dst)
neg_idx_list_dst = neg_idx_list_dst[: len(neg_idx_list_src)]
neg_idx_list_dst = torch.LongTensor(neg_idx_list_dst).to(device)
neg_dst_emb = batch_context_embed[neg_idx_list_dst]
else:
neg_dst = choice(
self.g.num_nodes(), size=len(neg_src_emb), prob=self.neg_prob
)
neg_dst_emb = self.context_embed(neg_dst.to(device))
pos_score = torch.sum(torch.mul(pos_src_emb, pos_dst_emb), dim=1)
pos_score = torch.clamp(pos_score, max=6, min=-6)
pos_score = torch.mean(-F.logsigmoid(pos_score))
neg_score = torch.sum(torch.mul(neg_src_emb, neg_dst_emb), dim=1)
neg_score = torch.clamp(neg_score, max=6, min=-6)
neg_score = (
torch.mean(-F.logsigmoid(-neg_score))
* self.negative_size
* self.neg_weight
)
return torch.mean(pos_score + neg_score)
[docs]class MetaPath2Vec(nn.Module):
r"""metapath2vec module from `metapath2vec: Scalable Representation Learning for
Heterogeneous Networks <https://dl.acm.org/doi/pdf/10.1145/3097983.3098036>`__
To achieve efficient optimization, we leverage the negative sampling technique for the
training process. Repeatedly for each node in meta-path, we treat it as the center node
and sample nearby positive nodes within context size and draw negative samples among all
types of nodes from all meta-paths. Then we can use the center-context paired nodes and
context-negative paired nodes to update the network.
Parameters
----------
g : DGLGraph
Graph for learning node embeddings. Two different canonical edge types
:attr:`(utype, etype, vtype)` are not allowed to have same :attr:`etype`.
metapath : list[str]
A sequence of edge types in the form of a string. It defines a new edge type by composing
multiple edge types in order. Note that the start node type and the end one are commonly
the same.
window_size : int
In a random walk :attr:`w`, a node :attr:`w[j]` is considered close to a node
:attr:`w[i]` if :attr:`i - window_size <= j <= i + window_size`.
emb_dim : int, optional
Size of each embedding vector. Default: 128
negative_size : int, optional
Number of negative samples to use for each positive sample. Default: 5
sparse : bool, optional
If True, gradients with respect to the learnable weights will be sparse.
Default: True
Attributes
----------
node_embed : nn.Embedding
Embedding table of all nodes
local_to_global_nid : dict[str, list]
Mapping from type-specific node IDs to global node IDs
Examples
--------
>>> import torch
>>> import dgl
>>> from torch.optim import SparseAdam
>>> from torch.utils.data import DataLoader
>>> from dgl.nn.pytorch import MetaPath2Vec
>>> # Define a model
>>> g = dgl.heterograph({
... ('user', 'uc', 'company'): dgl.rand_graph(100, 1000).edges(),
... ('company', 'cp', 'product'): dgl.rand_graph(100, 1000).edges(),
... ('company', 'cu', 'user'): dgl.rand_graph(100, 1000).edges(),
... ('product', 'pc', 'company'): dgl.rand_graph(100, 1000).edges()
... })
>>> model = MetaPath2Vec(g, ['uc', 'cu'], window_size=1)
>>> # Use the source node type of etype 'uc'
>>> dataloader = DataLoader(torch.arange(g.num_nodes('user')), batch_size=128,
... shuffle=True, collate_fn=model.sample)
>>> optimizer = SparseAdam(model.parameters(), lr=0.025)
>>> for (pos_u, pos_v, neg_v) in dataloader:
... loss = model(pos_u, pos_v, neg_v)
... optimizer.zero_grad()
... loss.backward()
... optimizer.step()
>>> # Get the embeddings of all user nodes
>>> user_nids = torch.LongTensor(model.local_to_global_nid['user'])
>>> user_emb = model.node_embed(user_nids)
"""
def __init__(
self,
g,
metapath,
window_size,
emb_dim=128,
negative_size=5,
sparse=True,
):
super().__init__()
assert (
len(metapath) + 1 >= window_size
), f"Expect len(metapath) >= window_size - 1, got {metapath} and {window_size}"
self.hg = g
self.emb_dim = emb_dim
self.metapath = metapath
self.window_size = window_size
self.negative_size = negative_size
# convert edge metapath to node metapath
# get initial source node type
src_type, _, _ = g.to_canonical_etype(metapath[0])
node_metapath = [src_type]
for etype in metapath:
_, _, dst_type = g.to_canonical_etype(etype)
node_metapath.append(dst_type)
self.node_metapath = node_metapath
# Convert the graph into a homogeneous one for global to local node ID mapping
g = to_homogeneous(g)
# Convert it back to the hetero one for local to global node ID mapping
hg = to_heterogeneous(g, self.hg.ntypes, self.hg.etypes)
local_to_global_nid = hg.ndata[NID]
for key, val in local_to_global_nid.items():
local_to_global_nid[key] = list(val.cpu().numpy())
self.local_to_global_nid = local_to_global_nid
num_nodes_total = hg.num_nodes()
node_frequency = torch.zeros(num_nodes_total)
# random walk
for idx in trange(hg.num_nodes(node_metapath[0])):
traces, _ = random_walk(g=hg, nodes=[idx], metapath=metapath)
for tr in traces.cpu().numpy():
tr_nids = [
self.local_to_global_nid[node_metapath[i]][tr[i]]
for i in range(len(tr))
]
node_frequency[torch.LongTensor(tr_nids)] += 1
neg_prob = node_frequency.pow(0.75)
self.neg_prob = neg_prob / neg_prob.sum()
# center node embedding
self.node_embed = nn.Embedding(
num_nodes_total, self.emb_dim, sparse=sparse
)
self.context_embed = nn.Embedding(
num_nodes_total, self.emb_dim, sparse=sparse
)
self.reset_parameters()
[docs] def reset_parameters(self):
"""Reinitialize learnable parameters"""
init_range = 1.0 / self.emb_dim
init.uniform_(self.node_embed.weight.data, -init_range, init_range)
init.constant_(self.context_embed.weight.data, 0)
def sample(self, indices):
"""Sample positive and negative samples
Parameters
----------
indices : torch.Tensor
Node IDs of the source node type from which we perform random walks
Returns
-------
torch.Tensor
Positive center nodes
torch.Tensor
Positive context nodes
torch.Tensor
Negative context nodes
"""
traces, _ = random_walk(
g=self.hg, nodes=indices, metapath=self.metapath
)
u_list = []
v_list = []
for tr in traces.cpu().numpy():
tr_nids = [
self.local_to_global_nid[self.node_metapath[i]][tr[i]]
for i in range(len(tr))
]
for i, u in enumerate(tr_nids):
for j, v in enumerate(
tr_nids[max(i - self.window_size, 0) : i + self.window_size]
):
if i == j:
continue
u_list.append(u)
v_list.append(v)
neg_v = choice(
self.hg.num_nodes(),
size=len(u_list) * self.negative_size,
prob=self.neg_prob,
).reshape(len(u_list), self.negative_size)
return torch.LongTensor(u_list), torch.LongTensor(v_list), neg_v
[docs] def forward(self, pos_u, pos_v, neg_v):
r"""Compute the loss for the batch of positive and negative samples
Parameters
----------
pos_u : torch.Tensor
Positive center nodes
pos_v : torch.Tensor
Positive context nodes
neg_v : torch.Tensor
Negative context nodes
Returns
-------
torch.Tensor
Loss value
"""
emb_u = self.node_embed(pos_u)
emb_v = self.context_embed(pos_v)
emb_neg_v = self.context_embed(neg_v)
score = torch.sum(torch.mul(emb_u, emb_v), dim=1)
score = torch.clamp(score, max=10, min=-10)
score = -F.logsigmoid(score)
neg_score = torch.bmm(emb_neg_v, emb_u.unsqueeze(2)).squeeze()
neg_score = torch.clamp(neg_score, max=10, min=-10)
neg_score = -torch.sum(F.logsigmoid(-neg_score), dim=1)
return torch.mean(score + neg_score)