dmlc--dgl
Minjie Wang
565f0c88fc
* refactor graph conv * docs & tests * fix lint * fix lint * fix lint * fix lint script * fix lint * Update * Style fix * Fix style * Fix style * Fix gpu case * Fix for gpu case * Hotfix edgesoftmax docs * Handle repeated features * Add docstring * Set default arguments * Remove dropout from nn.conv * Fix * add util fn for renaming * revert gcn_spmv.py * mx folder * fix wierd bug * fix mx * fix lint
84 行
2.5 KiB
Python
84 行
2.5 KiB
Python
"""GCN using builtin functions that enables SPMV optimization.
|
|
|
|
References:
|
|
- Semi-Supervised Classification with Graph Convolutional Networks
|
|
- Paper: https://arxiv.org/abs/1609.02907
|
|
- Code: https://github.com/tkipf/gcn
|
|
"""
|
|
import math
|
|
import torch
|
|
import torch.nn as nn
|
|
import dgl.function as fn
|
|
|
|
class GCNLayer(nn.Module):
|
|
def __init__(self,
|
|
g,
|
|
in_feats,
|
|
out_feats,
|
|
activation,
|
|
dropout,
|
|
bias=True):
|
|
super(GCNLayer, self).__init__()
|
|
self.g = g
|
|
self.weight = nn.Parameter(torch.Tensor(in_feats, out_feats))
|
|
if bias:
|
|
self.bias = nn.Parameter(torch.Tensor(out_feats))
|
|
else:
|
|
self.bias = None
|
|
self.activation = activation
|
|
if dropout:
|
|
self.dropout = nn.Dropout(p=dropout)
|
|
else:
|
|
self.dropout = 0.
|
|
self.reset_parameters()
|
|
|
|
def reset_parameters(self):
|
|
stdv = 1. / math.sqrt(self.weight.size(1))
|
|
self.weight.data.uniform_(-stdv, stdv)
|
|
if self.bias is not None:
|
|
self.bias.data.uniform_(-stdv, stdv)
|
|
|
|
def forward(self, h):
|
|
if self.dropout:
|
|
h = self.dropout(h)
|
|
h = torch.mm(h, self.weight)
|
|
# normalization by square root of src degree
|
|
h = h * self.g.ndata['norm']
|
|
self.g.ndata['h'] = h
|
|
self.g.update_all(fn.copy_src(src='h', out='m'),
|
|
fn.sum(msg='m', out='h'))
|
|
h = self.g.ndata.pop('h')
|
|
# normalization by square root of dst degree
|
|
h = h * self.g.ndata['norm']
|
|
# bias
|
|
if self.bias is not None:
|
|
h = h + self.bias
|
|
if self.activation:
|
|
h = self.activation(h)
|
|
return h
|
|
|
|
class GCN(nn.Module):
|
|
def __init__(self,
|
|
g,
|
|
in_feats,
|
|
n_hidden,
|
|
n_classes,
|
|
n_layers,
|
|
activation,
|
|
dropout):
|
|
super(GCN, self).__init__()
|
|
self.layers = nn.ModuleList()
|
|
# input layer
|
|
self.layers.append(GCNLayer(g, in_feats, n_hidden, activation, 0.))
|
|
# hidden layers
|
|
for i in range(n_layers - 1):
|
|
self.layers.append(GCNLayer(g, n_hidden, n_hidden, activation, dropout))
|
|
# output layer
|
|
self.layers.append(GCNLayer(g, n_hidden, n_classes, None, dropout))
|
|
|
|
def forward(self, features):
|
|
h = features
|
|
for layer in self.layers:
|
|
h = layer(h)
|
|
return h
|