Merge pull request tkipf#10 from daniilidis-group/master

Small changes to make it compatible with PyTorch 0.4.0

pygcn/layers.py

@@ -6,29 +6,6 @@
 from torch.nn.modules.module import Module
 
 
-class SparseMM(torch.autograd.Function):
-    """
-    Sparse x dense matrix multiplication with autograd support.
-
-    Implementation by Soumith Chintala:
-    https://discuss.pytorch.org/t/
-    does-pytorch-support-autograd-on-sparse-matrix/6156/7
-    """
-
-    def __init__(self, sparse):
-        super(SparseMM, self).__init__()
-        self.sparse = sparse
-
-    def forward(self, dense):
-        return torch.mm(self.sparse, dense)
-
-    def backward(self, grad_output):
-        grad_input = None
-        if self.needs_input_grad[0]:
-            grad_input = torch.mm(self.sparse.t(), grad_output)
-        return grad_input
-
-
 class GraphConvolution(Module):
     """
     Simple GCN layer, similar to https://arxiv.org/abs/1609.02907
@@ -38,9 +15,9 @@ def __init__(self, in_features, out_features, bias=True):
         super(GraphConvolution, self).__init__()
         self.in_features = in_features
         self.out_features = out_features
-        self.weight = Parameter(torch.Tensor(in_features, out_features))
+        self.weight = Parameter(torch.FloatTensor(in_features, out_features))
         if bias:
-            self.bias = Parameter(torch.Tensor(out_features))
+            self.bias = Parameter(torch.FloatTensor(out_features))
         else:
             self.register_parameter('bias', None)
         self.reset_parameters()
@@ -53,7 +30,7 @@ def reset_parameters(self):
 
     def forward(self, input, adj):
         support = torch.mm(input, self.weight)
-        output = SparseMM(adj)(support)
+        output = torch.spmm(adj, support)
         if self.bias is not None:
             return output + self.bias
         else:

pygcn/train.py

@@ -8,7 +8,6 @@
 import torch
 import torch.nn.functional as F
 import torch.optim as optim
-from torch.autograd import Variable
 
 from pygcn.utils import load_data, accuracy
 from pygcn.models import GCN
@@ -45,7 +44,7 @@
 # Model and optimizer
 model = GCN(nfeat=features.shape[1],
             nhid=args.hidden,
-            nclass=labels.max() + 1,
+            nclass=labels.max().item() + 1,
             dropout=args.dropout)
 optimizer = optim.Adam(model.parameters(),
                        lr=args.lr, weight_decay=args.weight_decay)
@@ -59,8 +58,6 @@
     idx_val = idx_val.cuda()
     idx_test = idx_test.cuda()
 
-features, labels = Variable(features), Variable(labels)
-
 
 def train(epoch):
     t = time.time()
@@ -94,8 +91,8 @@ def test():
     loss_test = F.nll_loss(output[idx_test], labels[idx_test])
     acc_test = accuracy(output[idx_test], labels[idx_test])
     print("Test set results:",
-          "loss= {:.4f}".format(loss_test.data[0]),
-          "accuracy= {:.4f}".format(acc_test.data[0]))
+          "loss= {:.4f}".format(loss_test.item()),
+          "accuracy= {:.4f}".format(acc_test.item()))
 
 
 # Train model