figure.png

    1. import  math
    3. import  torch
    4. from    torch import nn
    7. class GraphConvolution(nn.Module):
    8.     """
    9.     Simple GCN layer, similar to https://arxiv.org/abs/1609.02907
    10.     """
    12.     def __init__(self, in_features, out_features, bias=True):
    13.         super(GraphConvolution, self).__init__()
    15.         self.in_features = in_features
    16.         self.out_features = out_features
    18.         self.weight = nn.Parameter(torch.FloatTensor(in_features, out_features))
    19.         if bias:
    20.             self.bias = nn.Parameter(torch.FloatTensor(out_features))
    21.         else:
    22.             self.register_parameter('bias', None)
    23.         self.reset_parameters()
    25.     def reset_parameters(self):
    26.         stdv = 1. / math.sqrt(self.weight.size(1))
    27.         self.weight.data.uniform_(-stdv, stdv)
    28.         if self.bias is not None:
    29.             self.bias.data.uniform_(-stdv, stdv)
    31.     def forward(self, input, adj):
    32.         """
    34.         :param input:
    35.         :param adj:
    36.         :return:
    37.         """
    38.         support = torch.mm(input, self.weight)
    39.         # sparse matrix multiplication
    40.         output = torch.spmm(adj, support)
    41.         if self.bias is not None:
    42.             return output + self.bias
    43.         else:
    44.             return output
    46.     def __repr__(self):
    47.         return self.__class__.__name__ + ' (' \
    48.                + str(self.in_features) + ' -> ' \
    49.                + str(self.out_features) + ')'
    1. import  torch.nn as nn
    2. import  torch.nn.functional as F
    3. from    layers import GraphConvolution
    6. class GCN(nn.Module):
    8.     def __init__(self, nfeat, nhid, nclass, dropout):
    9.         super(GCN, self).__init__()
    11.         self.gc1 = GraphConvolution(nfeat, nhid)
    12.         self.gc2 = GraphConvolution(nhid, nclass)
    13.         self.dropout = dropout
    15.     def forward(self, x, adj):
    16.         """
    18.         :param x: [2708, 1433]
    19.         :param adj: [2708, 2708]
    20.         :return:
    21.         """
    22.         # print('x:', x.shape, 'adj:', adj.shape)
    23.         # => [2708, 16]
    24.         x = F.relu(self.gc1(x, adj))
    25.         # print('gcn1:', x.shape)
    26.         x = F.dropout(x, self.dropout, training=self.training)
    27.         # => [2708, 7]
    28.         x = self.gc2(x, adj)
    29.         # print('gcn2:', x.shape)
    30.         return F.log_softmax(x, dim=1)
    1. from __future__ import division
    2. from __future__ import print_function
    4. import time
    5. import argparse
    6. import numpy as np
    8. import torch
    9. import torch.nn.functional as F
    10. import torch.optim as optim
    12. from utils import load_data, accuracy
    13. from models import GCN
    15. # Training settings
    16. parser = argparse.ArgumentParser()
    17. parser.add_argument('--no-cuda', action='store_true', default=False,
    18.                     help='Disables CUDA training.')
    19. parser.add_argument('--fastmode', action='store_true', default=False,
    20.                     help='Validate during training pass.')
    21. parser.add_argument('--seed', type=int, default=42, help='Random seed.')
    22. parser.add_argument('--epochs', type=int, default=200,
    23.                     help='Number of epochs to train.')
    24. parser.add_argument('--lr', type=float, default=0.01,
    25.                     help='Initial learning rate.')
    26. parser.add_argument('--weight_decay', type=float, default=5e-4,
    27.                     help='Weight decay (L2 loss on parameters).')
    28. parser.add_argument('--hidden', type=int, default=16,
    29.                     help='Number of hidden units.')
    30. parser.add_argument('--dropout', type=float, default=0.5,
    31.                     help='Dropout rate (1 - keep probability).')
    33. args = parser.parse_args()
    34. args.cuda = not args.no_cuda and torch.cuda.is_available()
    36. np.random.seed(args.seed)
    37. torch.manual_seed(args.seed)
    38. if args.cuda:
    39.     torch.cuda.manual_seed(args.seed)
    41. # Load data
    42. adj, features, labels, idx_train, idx_val, idx_test = load_data()
    44. # Model and optimizer
    45. model = GCN(nfeat=features.shape[1],
    46.             nhid=args.hidden,
    47.             nclass=labels.max().item() + 1,
    48.             dropout=args.dropout)
    49. optimizer = optim.Adam(model.parameters(),
    50.                        lr=args.lr, weight_decay=args.weight_decay)
    52. print(model)
    54. if args.cuda:
    55.     model.cuda()
    56.     features = features.cuda()
    57.     adj = adj.cuda()
    58.     labels = labels.cuda()
    59.     idx_train = idx_train.cuda()
    60.     idx_val = idx_val.cuda()
    61.     idx_test = idx_test.cuda()
    64. def train(epoch):
    66.     t = time.time()
    68.     model.train()
    69.     # [N, 7]
    70.     output = model(features, adj)
    71.     loss_train = F.nll_loss(output[idx_train], labels[idx_train])
    72.     acc_train = accuracy(output[idx_train], labels[idx_train])
    74.     optimizer.zero_grad()
    75.     loss_train.backward()
    76.     optimizer.step()
    78.     if not args.fastmode:
    79.         # Evaluate validation set performance separately,
    80.         # deactivates dropout during validation run.
    81.         model.eval()
    82.         output = model(features, adj)
    84.     loss_val = F.nll_loss(output[idx_val], labels[idx_val])
    85.     acc_val = accuracy(output[idx_val], labels[idx_val])
    86.     print('Epoch: {:04d}'.format(epoch+1),
    87.           'loss_train: {:.4f}'.format(loss_train.item()),
    88.           'acc_train: {:.4f}'.format(acc_train.item()),
    89.           'loss_val: {:.4f}'.format(loss_val.item()),
    90.           'acc_val: {:.4f}'.format(acc_val.item()),
    91.           'time: {:.4f}s'.format(time.time() - t))
    94. def test():
    96.     model.eval()
    98.     output = model(features, adj)
    99.     loss_test = F.nll_loss(output[idx_test], labels[idx_test])
    100.     acc_test = accuracy(output[idx_test], labels[idx_test])
    101.     print("Test set results:",
    102.           "loss= {:.4f}".format(loss_test.item()),
    103.           "accuracy= {:.4f}".format(acc_test.item()))
    106. # Train model
    107. t_total = time.time()
    108. for epoch in range(args.epochs):
    109.     train(epoch)
    110. print("Optimization Finished!")
    111. print("Total time elapsed: {:.4f}s".format(time.time() - t_total))
    113. # Testing
    114. test()
    1. import numpy as np
    2. import scipy.sparse as sp
    3. import torch
    6. def encode_onehot(labels):
    7.     classes = set(labels)
    8.     classes_dict = {c: np.identity(len(classes))[i, :] for i, c in
    9.                     enumerate(classes)}
    10.     labels_onehot = np.array(list(map(classes_dict.get, labels)),
    11.                              dtype=np.int32)
    12.     return labels_onehot
    15. def load_data(path="./data/cora/", dataset="cora"):
    16.     """Load citation network dataset (cora only for now)"""
    17.     print('Loading {} dataset...'.format(dataset))
    19.     idx_features_labels = np.genfromtxt("{}{}.content".format(path, dataset),
    20.                                         dtype=np.dtype(str))
    21.     features = sp.csr_matrix(idx_features_labels[:, 1:-1], dtype=np.float32)
    22.     labels = encode_onehot(idx_features_labels[:, -1])
    24.     # build graph
    25.     idx = np.array(idx_features_labels[:, 0], dtype=np.int32)
    26.     idx_map = {j: i for i, j in enumerate(idx)}
    27.     edges_unordered = np.genfromtxt("{}{}.cites".format(path, dataset),
    28.                                     dtype=np.int32)
    29.     edges = np.array(list(map(idx_map.get, edges_unordered.flatten())),
    30.                      dtype=np.int32).reshape(edges_unordered.shape)
    31.     adj = sp.coo_matrix((np.ones(edges.shape[0]), (edges[:, 0], edges[:, 1])),
    32.                         shape=(labels.shape[0], labels.shape[0]),
    33.                         dtype=np.float32)
    35.     # build symmetric adjacency matrix
    36.     adj = adj + adj.T.multiply(adj.T > adj) - adj.multiply(adj.T > adj)
    38.     features = normalize(features)
    39.     adj = normalize(adj + sp.eye(adj.shape[0]))
    41.     idx_train = range(140)
    42.     idx_val = range(200, 500)
    43.     idx_test = range(500, 1500)
    45.     features = torch.FloatTensor(np.array(features.todense()))
    46.     labels = torch.LongTensor(np.where(labels)[1])
    47.     adj = sparse_mx_to_torch_sparse_tensor(adj)
    49.     idx_train = torch.LongTensor(idx_train)
    50.     idx_val = torch.LongTensor(idx_val)
    51.     idx_test = torch.LongTensor(idx_test)
    53.     return adj, features, labels, idx_train, idx_val, idx_test
    56. def normalize(mx):
    57.     """Row-normalize sparse matrix"""
    58.     rowsum = np.array(mx.sum(1))
    59.     r_inv = np.power(rowsum, -1).flatten()
    60.     r_inv[np.isinf(r_inv)] = 0.
    61.     r_mat_inv = sp.diags(r_inv)
    62.     mx = r_mat_inv.dot(mx)
    63.     return mx
    66. def accuracy(output, labels):
    67.     preds = output.max(1)[1].type_as(labels)
    68.     correct = preds.eq(labels).double()
    69.     correct = correct.sum()
    70.     return correct / len(labels)
    73. def sparse_mx_to_torch_sparse_tensor(sparse_mx):
    74.     """Convert a scipy sparse matrix to a torch sparse tensor."""
    75.     sparse_mx = sparse_mx.tocoo().astype(np.float32)
    76.     indices = torch.from_numpy(
    77.         np.vstack((sparse_mx.row, sparse_mx.col)).astype(np.int64))
    78.     values = torch.from_numpy(sparse_mx.data)
    79.     shape = torch.Size(sparse_mx.shape)
    80.     return torch.sparse.FloatTensor(indices, values, shape)