Skip to content

GitLab

Commit 378c2645 authored by Lingfan Yu's avatar Lingfan Yu Committed by Minjie Wang
Browse files

[Model] Fix GCN Normalization (#249) 

* WIP

* lr -> 0.01

* new cora dataset

* normalization code

* minor format change

* normalization factor for deg bucket
parent 21255b65
Show whitespace changes
Inline Side-by-side
Showing with 183 additions and 40 deletions
examples/pytorch/gcn/gcn.py
View file @ 378c2645
......@@ -27,7 +27,9 @@ class NodeApplyModule(nn.Module):
self.activation = activation
def forward(self, nodes):
h = self.linear(nodes.data['h'])
# normalization by square root of dst degree
h = nodes.data['h'] * nodes.data['norm']
h = self.linear(h)
if self.activation:
h = self.activation(h)
return {'h' : h}
......@@ -49,8 +51,10 @@ class GCN(nn.Module):
else:
self.dropout = 0.
self.layers = nn.ModuleList()
# input layer
self.layers = nn.ModuleList([NodeApplyModule(in_feats, n_hidden, activation)])
self.layers.append(NodeApplyModule(in_feats, n_hidden, activation))
# hidden layers
for i in range(n_layers - 1):
......@@ -62,22 +66,34 @@ class GCN(nn.Module):
def forward(self, features):
self.g.ndata['h'] = features
for layer in self.layers:
for idx, layer in enumerate(self.layers):
# apply dropout
if self.dropout:
self.g.apply_nodes(apply_node_func=
lambda nodes: {'h': self.dropout(nodes.data['h'])})
if idx > 0 and self.dropout:
self.g.ndata['h'] = self.dropout(self.g.ndata['h'])
# normalization by square root of src degree
self.g.ndata['h'] = self.g.ndata['h'] * self.g.ndata['norm']
self.g.update_all(gcn_msg, gcn_reduce, layer)
return self.g.ndata.pop('h')
def evaluate(model, features, labels, mask):
model.eval()
with torch.no_grad():
logits = model(features)
logits = logits[mask]
labels = labels[mask]
_, indices = torch.max(logits, dim=1)
correct = torch.sum(indices == labels)
return correct.item() * 1.0 / len(labels)
def main(args):
# load and preprocess dataset
# Todo: adjacency normalization
data = load_data(args)
features = torch.FloatTensor(data.features)
labels = torch.LongTensor(data.labels)
mask = torch.ByteTensor(data.train_mask)
train_mask = torch.ByteTensor(data.train_mask)
val_mask = torch.ByteTensor(data.val_mask)
test_mask = torch.ByteTensor(data.test_mask)
in_feats = features.shape[1]
n_classes = data.num_labels
n_edges = data.graph.number_of_edges()
......@@ -89,10 +105,24 @@ def main(args):
torch.cuda.set_device(args.gpu)
features = features.cuda()
labels = labels.cuda()
mask = mask.cuda()
train_mask = train_mask.cuda()
val_mask = val_mask.cuda()
test_mask = test_mask.cuda()
# create GCN model
# graph preprocess and calculate normalization factor
g = DGLGraph(data.graph)
n_edges = g.number_of_edges()
# add self loop
g.add_edges(g.nodes(), g.nodes())
# normalization
degs = g.in_degrees().float()
norm = torch.pow(degs, -0.5)
norm[torch.isinf(norm)] = 0
if cuda:
norm = norm.cuda()
g.ndata['norm'] = norm.unsqueeze(1)
# create GCN model
model = GCN(g,
in_feats,
args.n_hidden,
......@@ -105,17 +135,20 @@ def main(args):
model.cuda()
# use optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
# initialize graph
dur = []
for epoch in range(args.n_epochs):
model.train()
if epoch >= 3:
t0 = time.time()
# forward
logits = model(features)
logp = F.log_softmax(logits, 1)
loss = F.nll_loss(logp[mask], labels[mask])
loss = F.nll_loss(logp[train_mask], labels[train_mask])
optimizer.zero_grad()
loss.backward()
......@@ -124,24 +157,33 @@ def main(args):
if epoch >= 3:
dur.append(time.time() - t0)
print("Epoch {:05d} | Loss {:.4f} | Time(s) {:.4f} | ETputs(KTEPS) {:.2f}".format(
epoch, loss.item(), np.mean(dur), n_edges / np.mean(dur) / 1000))
acc = evaluate(model, features, labels, val_mask)
print("Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | Accuracy {:.4f} | "
"ETputs(KTEPS) {:.2f}".format(epoch, np.mean(dur), loss.item(),
acc, n_edges / np.mean(dur) / 1000))
print()
acc = evaluate(model, features, labels, test_mask)
print("Test Accuracy {:.4f}".format(acc))
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='GCN')
register_data_args(parser)
parser.add_argument("--dropout", type=float, default=0,
parser.add_argument("--dropout", type=float, default=0.5,
help="dropout probability")
parser.add_argument("--gpu", type=int, default=-1,
help="gpu")
parser.add_argument("--lr", type=float, default=1e-3,
parser.add_argument("--lr", type=float, default=1e-2,
help="learning rate")
parser.add_argument("--n-epochs", type=int, default=20,
parser.add_argument("--n-epochs", type=int, default=200,
help="number of training epochs")
parser.add_argument("--n-hidden", type=int, default=16,
help="number of hidden gcn units")
parser.add_argument("--n-layers", type=int, default=1,
help="number of hidden gcn layers")
parser.add_argument("--weight-decay", type=float, default=5e-4,
help="Weight for L2 loss")
args = parser.parse_args()
print(args)
......
examples/pytorch/gcn/gcn_spmv.py
View file @ 378c2645
......@@ -18,15 +18,16 @@ from dgl.data import register_data_args, load_data
class NodeApplyModule(nn.Module):
def __init__(self, in_feats, out_feats, activation=None):
super(NodeApplyModule, self).__init__()
self.linear = nn.Linear(in_feats, out_feats)
nn.init.xavier_normal_(self.linear.weight)
self.activation = activation
def forward(self, nodes):
h = self.linear(nodes.data['h'])
# normalization by square root of dst degree
h = nodes.data['h'] * nodes.data['norm']
h = self.linear(h)
if self.activation:
h = self.activation(h)
return {'h': h}
class GCN(nn.Module):
......@@ -46,8 +47,10 @@ class GCN(nn.Module):
else:
self.dropout = 0.
self.layers = nn.ModuleList()
# input layer
self.layers = nn.ModuleList([NodeApplyModule(in_feats, n_hidden, activation)])
self.layers.append(NodeApplyModule(in_feats, n_hidden, activation))
# hidden layers
for i in range(n_layers - 1):
......@@ -59,24 +62,35 @@ class GCN(nn.Module):
def forward(self, features):
self.g.ndata['h'] = features
for layer in self.layers:
for idx, layer in enumerate(self.layers):
# apply dropout
if self.dropout:
self.g.apply_nodes(apply_node_func=
lambda nodes: {'h': self.dropout(nodes.data['h'])})
if idx > 0 and self.dropout:
self.g.ndata['h'] = self.dropout(self.g.ndata['h'])
# normalization by square root of src degree
self.g.ndata['h'] = self.g.ndata['h'] * self.g.ndata['norm']
self.g.update_all(fn.copy_src(src='h', out='m'),
fn.sum(msg='m', out='h'),
layer)
return self.g.pop_n_repr('h')
def evaluate(model, features, labels, mask):
model.eval()
with torch.no_grad():
logits = model(features)
logits = logits[mask]
labels = labels[mask]
_, indices = torch.max(logits, dim=1)
correct = torch.sum(indices == labels)
return correct.item() * 1.0 / len(labels)
def main(args):
# load and preprocess dataset
# Todo: adjacency normalization
data = load_data(args)
features = torch.FloatTensor(data.features)
labels = torch.LongTensor(data.labels)
mask = torch.ByteTensor(data.train_mask)
train_mask = torch.ByteTensor(data.train_mask)
val_mask = torch.ByteTensor(data.val_mask)
test_mask = torch.ByteTensor(data.test_mask)
in_feats = features.shape[1]
n_classes = data.num_labels
n_edges = data.graph.number_of_edges()
......@@ -88,10 +102,24 @@ def main(args):
torch.cuda.set_device(args.gpu)
features = features.cuda()
labels = labels.cuda()
mask = mask.cuda()
train_mask = train_mask.cuda()
val_mask = val_mask.cuda()
test_mask = test_mask.cuda()
# create GCN model
# graph preprocess and calculate normalization factor
g = DGLGraph(data.graph)
n_edges = g.number_of_edges()
# add self loop
g.add_edges(g.nodes(), g.nodes())
# normalization
degs = g.in_degrees().float()
norm = torch.pow(degs, -0.5)
norm[torch.isinf(norm)] = 0
if cuda:
norm = norm.cuda()
g.ndata['norm'] = norm.unsqueeze(1)
# create GCN model
model = GCN(g,
in_feats,
args.n_hidden,
......@@ -104,17 +132,20 @@ def main(args):
model.cuda()
# use optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
# initialize graph
dur = []
for epoch in range(args.n_epochs):
model.train()
if epoch >= 3:
t0 = time.time()
# forward
logits = model(features)
logp = F.log_softmax(logits, 1)
loss = F.nll_loss(logp[mask], labels[mask])
loss = F.nll_loss(logp[train_mask], labels[train_mask])
optimizer.zero_grad()
loss.backward()
......@@ -123,24 +154,33 @@ def main(args):
if epoch >= 3:
dur.append(time.time() - t0)
print("Epoch {:05d} | Loss {:.4f} | Time(s) {:.4f} | ETputs(KTEPS) {:.2f}".format(
epoch, loss.item(), np.mean(dur), n_edges / np.mean(dur) / 1000))
acc = evaluate(model, features, labels, val_mask)
print("Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | Accuracy {:.4f} | "
"ETputs(KTEPS) {:.2f}". format(epoch, np.mean(dur), loss.item(),
acc, n_edges / np.mean(dur) / 1000))
print()
acc = evaluate(model, features, labels, test_mask)
print("Test Accuracy {:.4f}".format(acc))
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='GCN')
register_data_args(parser)
parser.add_argument("--dropout", type=float, default=0,
parser.add_argument("--dropout", type=float, default=0.5,
help="dropout probability")
parser.add_argument("--gpu", type=int, default=-1,
help="gpu")
parser.add_argument("--lr", type=float, default=1e-3,
parser.add_argument("--lr", type=float, default=1e-2,
help="learning rate")
parser.add_argument("--n-epochs", type=int, default=20,
parser.add_argument("--n-epochs", type=int, default=200,
help="number of training epochs")
parser.add_argument("--n-hidden", type=int, default=16,
help="number of hidden gcn units")
parser.add_argument("--n-layers", type=int, default=1,
help="number of hidden gcn layers")
parser.add_argument("--weight-decay", type=float, default=5e-4,
help="Weight for L2 loss")
args = parser.parse_args()
print(args)
......
python/dgl/data/citation_graph.py
View file @ 378c2645
......@@ -15,7 +15,7 @@ import dgl
from .utils import download, extract_archive, get_download_dir, _get_dgl_url
_urls = {
'cora' : 'dataset/cora.zip',
'cora' : 'dataset/cora_raw.zip',
'citeseer' : 'dataset/citeseer.zip',
'pubmed' : 'dataset/pubmed.zip',
'cora_binary' : 'dataset/cora_binary.zip',
......@@ -140,7 +140,7 @@ def _sample_mask(idx, l):
return mask
def load_cora():
data = CitationGraphDataset('cora')
data = CoraDataset()
return data
def load_citeseer():
......@@ -335,3 +335,64 @@ class CoraBinary(object):
batched_pmpds = sp.block_diag(pmpds)
batched_labels = np.concatenate(labels, axis=0)
return batched_graphs, batched_pmpds, batched_labels
class CoraDataset(object):
def __init__(self):
self.name = 'cora'
self.dir = get_download_dir()
self.zip_file_path='{}/{}.zip'.format(self.dir, self.name)
download(_get_dgl_url(_urls[self.name]), path=self.zip_file_path)
extract_archive(self.zip_file_path,
'{}/{}'.format(self.dir, self.name))
self._load()
def _load(self):
idx_features_labels = np.genfromtxt("{}/cora/cora.content".
format(self.dir),
dtype=np.dtype(str))
features = sp.csr_matrix(idx_features_labels[:, 1:-1],
dtype=np.float32)
labels = _encode_onehot(idx_features_labels[:, -1])
self.num_labels = labels.shape[1]
# build graph
idx = np.array(idx_features_labels[:, 0], dtype=np.int32)
idx_map = {j: i for i, j in enumerate(idx)}
edges_unordered = np.genfromtxt("{}/cora/cora.cites".format(self.dir),
dtype=np.int32)
edges = np.array(list(map(idx_map.get, edges_unordered.flatten())),
dtype=np.int32).reshape(edges_unordered.shape)
adj = sp.coo_matrix((np.ones(edges.shape[0]),
(edges[:, 0], edges[:, 1])),
shape=(labels.shape[0], labels.shape[0]),
dtype=np.float32)
# build symmetric adjacency matrix
adj = adj + adj.T.multiply(adj.T > adj) - adj.multiply(adj.T > adj)
self.graph = nx.from_scipy_sparse_matrix(adj, create_using=nx.DiGraph())
features = _normalize(features)
self.features = np.array(features.todense())
self.labels = np.where(labels)[1]
self.train_mask = _sample_mask(range(140), labels.shape[0])
self.val_mask = _sample_mask(range(200, 500), labels.shape[0])
self.test_mask = _sample_mask(range(500, 1500), labels.shape[0])
def _normalize(mx):
"""Row-normalize sparse matrix"""
rowsum = np.array(mx.sum(1))
r_inv = np.power(rowsum, -1).flatten()
r_inv[np.isinf(r_inv)] = np.inf
r_mat_inv = sp.diags(r_inv)
mx = r_mat_inv.dot(mx)
return mx
def _encode_onehot(labels):
classes = set(labels)
classes_dict = {c: np.identity(len(classes))[i, :] for i, c in
enumerate(classes)}
labels_onehot = np.array(list(map(classes_dict.get, labels)),
dtype=np.int32)
return labels_onehot
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment