kaolin.ops.gcn¶
API¶
-
class
kaolin.ops.gcn.
GraphConv
(input_dim, output_dim, self_layer=True, bias=True)¶ Bases:
torch.nn.modules.module.Module
A simple graph convolution layer, similar to the one defined in Kipf et al. https://arxiv.org/abs/1609.02907
This operation with self_layer=False is equivalent to \((A H W)\) where: - \(H\) is the node features with shape (batch_size, num_nodes, input_dim) - \(W\) is a weight matrix of shape (input_dim, output_dim) - \(A\) is the adjacency matrix of shape (num_nodes, num_nodes). It can include self-loop.
With normalize_adj=True, it is equivalent to \((D^{-1} A H W)\), where: - \(D\) is a diagonal matrix with \(D_{ii}\) = the sum of the i-th row of A. In other words, \(D\) is the incoming degree of each node.
With self_layer=True, it is equivalent to the above plus \((H W_{\text{self}})\), where: - \(W_{\text{self}}\) is a separate weight matrix to filter each node’s self features.
Note that when self_layer is True, A should not include self-loop.
Example
>>> node_feat = torch.rand(1, 3, 5) >>> i = torch.LongTensor( ... [[0, 1, 1, 2, 2, 0], [1, 0, 2, 1, 0, 2]]) >>> v = torch.FloatTensor([1, 1, 1, 1, 1, 1]) >>> adj = torch.sparse.FloatTensor(i, v, torch.Size([3, 3])) >>> model = GraphConv(5, 10) >>> output = model(node_feat, adj) >>> # pre-normalize adj >>> adj = normalize_adj(adj) >>> output = model(node_feat, adj, normalize_adj=False)
If you use this code, please cite the original paper in addition to Kaolin. .. code-block:
@article{kipf2016semi, title={Semi-Supervised Classification with Graph Convolutional Networks}, author={Kipf, Thomas N and Welling, Max}, journal={arXiv preprint arXiv:1609.02907}, year={2016} }
Parameters: -
forward
(node_feat, adj, normalize_adj=True)¶ Parameters: - node_feat (torch.FloatTensor) – Shape = (batch_size, num_nodes, input_dim) The input features of each node.
- adj (torch.sparse.FloatTensor or torch.FloatTensor) – Shape = (num_nodes, num_nodes) The adjacency matrix. adj[i, j] is non-zero if there’s an incoming edge from j to i. Should not include self-loop if self_layer is True.
- normalize_adj (bool) – Set this to true to apply normalization to adjacency; that is, each output feature will be divided by the number of incoming neighbors. If normalization is not desired, or if the adjacency matrix is pre-normalized, set this to False to improve performance.
Returns: The output features of each node. Shape = (batch_size, num_nodes, output_dim)
Return type: (torch.FloatTensor)
-
initialize
()¶
-
-
kaolin.ops.gcn.
normalize_adj
(adj)¶ Normalize the adjacency matrix with shape = (num_nodes, num_nodes) such that the sum of each row is 1.
This operation is slow, so it should be done only once for a graph and then reused.
This supports both sparse tensor and regular tensor. The return type will be the same as the input type. For example, if the input is a sparse tensor, the normalized matrix will also be a sparse tensor.
Parameters: adj (torch.sparse.FloatTensor or torch.FloatTensor) – Shape = (num_nodes, num_nodes) The adjacency matrix. Returns: A new adjacency matrix with the same connectivity as the input, but with the sum of each row normalized to 1. Return type: (torch.sparse.FloatTensor or torch.FloatTensor)
-
kaolin.ops.gcn.
sparse_bmm
(sparse_matrix, dense_matrix_batch)¶ Perform torch.bmm on an unbatched sparse matrix and a batched dense matrix.
Parameters: - sparse_matrix (torch.sparse.FloatTensor) – Shape = (m, n)
- dense_matrix_batch (torch.FloatTensor) – Shape = (b, n, p)
Returns: Result of the batched matrix multiplication. Shape = (b, n, p)
Return type: (torch.FloatTensor)