csrmm2 not supported anymore, replaced with SpMM()

Author: alokpathy

.gitignore

@@ -1,28 +1,12 @@
 __pycache__/
+data/
 job_scripts/
 outputs/
 profs/
-data/
-.*.swm
-.*.swn
-.*.swo
+sparse-extension/build/
+sparse-extension/dist/
+sparse-extension/sparse_coo_tensor_cpp.egg-info/
 .*.swp
 .nfs*
-test*
-distr*.txt
-distr*.prof
-gcn*.txt
-*saved*
-gcn.py
-gcn_conv.py
-gcn_distr_cpu.py
-gcn_distr_1_5d_gpu_edit.py
-gcn_distr_2d_actandnorm.py
-gcn_distr_2d_gen.py
-gcn_distr_2d_cpu.py
-gcn_distr_2d_gpu.py
-gcn_distr_2d_gpu_gen.py
-gcn_distr_2d_gpu_floor_edit.py
-gcn_distr_2d_gpu_floor_old.py
-gcn_distr_2d_gpu_floor_gpu_gen.py
+*.txt
 *.pt

gcn.py

@@ -0,0 +1,126 @@
+import os.path as osp
+import argparse
+
+import torch
+from torch.nn import Parameter
+import torch.nn.functional as F
+from torch_geometric.datasets import Planetoid, PPI, Reddit
+import torch_geometric.transforms as T
+from torch_geometric.nn import GCNConv # noqa
+
+import time
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--use_gdc', action='store_true',
+                    help='Use GDC preprocessing.')
+args = parser.parse_args()
+
+dataset = 'Cora'
+path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', dataset)
+dataset = Planetoid(path, dataset, transform=T.NormalizeFeatures())
+# dataset = PPI(path, 'train', T.NormalizeFeatures())
+# dataset = Reddit(path, T.NormalizeFeatures())
+# dataset = Yelp(path, T.NormalizeFeatures())
+data = dataset[0]
+
+seed = 0
+
+if args.use_gdc:
+    gdc = T.GDC(self_loop_weight=1, normalization_in='sym',
+                normalization_out='col',
+                diffusion_kwargs=dict(method='ppr', alpha=0.05),
+                sparsification_kwargs=dict(method='topk', k=128,
+                                           dim=0), exact=True)
+    data = gdc(data)
+
+
+class Net(torch.nn.Module):
+    def __init__(self):
+        super(Net, self).__init__()
+        self.conv1 = GCNConv(dataset.num_features, 16, cached=True, normalize=False, bias=False)
+        self.conv2 = GCNConv(16, dataset.num_classes, cached=True, normalize=False, bias=False)
+
+        self.conv1.node_dim = 0
+        self.conv2.node_dim = 0
+
+        with torch.no_grad():
+            self.conv1.weight = Parameter(weight1)
+            self.conv2.weight = Parameter(weight2)
+        # self.conv1 = ChebConv(data.num_features, 16, K=2)
+        # self.conv2 = ChebConv(16, data.num_features, K=2)
+
+    def forward(self):
+        x, edge_index = data.x, data.edge_index
+        # x = F.relu(self.conv1(x, edge_index))
+        # x = F.dropout(x, training=self.training)
+        # x = self.conv2(x, edge_index)
+        # return F.log_softmax(x, dim=1)
+        x = self.conv1(x, edge_index)
+        # x = F.relu(x)
+        x = self.conv2(x, edge_index)
+        # return F.log_softmax(x, dim=1)
+        return x
+
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+torch.manual_seed(seed)
+weight1 = torch.rand(dataset.num_features, 16)
+weight1 = weight1.to(device)
+
+weight2 = torch.rand(16, dataset.num_classes)
+weight2 = weight2.to(device)
+
+data.y = data.y.type(torch.LongTensor)
+model, data = Net().to(device), data.to(device)
+
+optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+
+def train():
+    model.train()
+    optimizer.zero_grad()
+    outputs = model()
+    
+    # Note: bool type removes warnings, unsure of perf penalty
+    F.nll_loss(outputs[data.train_mask.bool()], data.y[data.train_mask.bool()]).backward()
+    # F.nll_loss(outputs, torch.max(data.y, 1)[1]).backward()
+
+    for W in model.parameters():
+        if W.grad is not None:
+            print(W.grad)
+
+    optimizer.step()
+    return outputs
+
+def test(outputs):
+    model.eval()
+    logits, accs = outputs, []
+    for _, mask in data('train_mask', 'val_mask', 'test_mask'):
+        pred = logits[mask].max(1)[1]
+        acc = pred.eq(data.y[mask]).sum().item() / mask.sum().item()
+        accs.append(acc)
+    return accs
+
+def main(): 
+    best_val_acc = test_acc = 0
+    outputs = None
+
+    tstart = time.time()
+
+    # for epoch in range(1, 101):
+    for epoch in range(1):
+        outputs = train()
+        train_acc, val_acc, tmp_test_acc = test(outputs)
+        if val_acc > best_val_acc:
+            best_val_acc = val_acc
+            test_acc = tmp_test_acc
+        log = 'Epoch: {:03d}, Train: {:.4f}, Val: {:.4f}, Test: {:.4f}'
+        print(log.format(epoch, train_acc, best_val_acc, test_acc))
+
+    tstop = time.time()
+    print("Time: " + str(tstop - tstart))
+
+    return outputs
+
+if __name__=='__main__':
+    print(main())

gcn_distr.py

@@ -235,9 +235,10 @@ def broad_func(node_count, am_partitions, inputs, rank, size, group):
 
         tstart_comp = start_time(group, rank)
 
-        spmm_gpu(am_partitions[i].indices()[0].int(), am_partitions[i].indices()[1].int(), 
-                        am_partitions[i].values(), am_partitions[i].size(0), 
-                        am_partitions[i].size(1), inputs_recv, z_loc)
+        # spmm_gpu(am_partitions[i].indices()[0].int(), am_partitions[i].indices()[1].int(), 
+        #                 am_partitions[i].values(), am_partitions[i].size(0), 
+        #                 am_partitions[i].size(1), inputs_recv, z_loc)
+        z_loc += torch.mm(am_partitions[i], inputs_recv)
 
         dur = stop_time(group, rank, tstart_comp)
         comp_time[run][rank] += dur
@@ -472,6 +473,9 @@ def oned_partition(rank, size, inputs, adj_matrix, data, features, classes, devi
     am_partitions = None
     am_pbyp = None
 
+    inputs = inputs.to(torch.device("cpu"))
+    adj_matrix = adj_matrix.to(torch.device("cpu"))
+
     # Compute the adj_matrix and inputs partitions for this process
     # TODO: Maybe I do want grad here. Unsure.
     with torch.no_grad():
@@ -530,7 +534,6 @@ def run(rank, size, inputs, adj_matrix, data, features, classes, device):
     # adj_matrix_loc = torch.rand(node_count, n_per_proc)
     # inputs_loc = torch.rand(n_per_proc, inputs.size(1))
 
-
     inputs_loc, adj_matrix_loc, am_pbyp = oned_partition(rank, size, inputs, adj_matrix, data, 
                                                                 features, classes, device)
 
@@ -682,6 +685,7 @@ def main():
         outputs = None
         if "OMPI_COMM_WORLD_RANK" in os.environ.keys():
             os.environ["RANK"] = os.environ["OMPI_COMM_WORLD_RANK"]
+
         # Initialize distributed environment with SLURM
         if "SLURM_PROCID" in os.environ.keys():
             os.environ["RANK"] = os.environ["SLURM_PROCID"]
@@ -701,14 +705,15 @@ def main():
         # device = torch.device('cpu')
         devid = rank_to_devid(rank, acc_per_rank)
         device = torch.device('cuda:{}'.format(devid))
+        print(f"device: {device}")
         torch.cuda.set_device(device)
         curr_devid = torch.cuda.current_device()
         # print(f"curr_devid: {curr_devid}", flush=True)
         devcount = torch.cuda.device_count()
 
     if graphname == "Cora":
         path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', graphname)
-        dataset = Planetoid(path, graphname, T.NormalizeFeatures())
+        dataset = Planetoid(path, graphname, transform=T.NormalizeFeatures())
         data = dataset[0]
         data = data.to(device)
         data.x.requires_grad = True

gcn_distr_15d.py

@@ -502,6 +502,9 @@ def oned_partition(rank, size, inputs, adj_matrix, data, features, classes, devi
     am_partitions = None
     am_pbyp = None
 
+    inputs = inputs.to(torch.device("cpu"))
+    adj_matrix = adj_matrix.to(torch.device("cpu"))
+
     rank_c = rank // replication
     # Compute the adj_matrix and inputs partitions for this process
     # TODO: Maybe I do want grad here. Unsure.
@@ -723,6 +726,17 @@ def main():
         if "OMPI_COMM_WORLD_RANK" in os.environ.keys():
             os.environ["RANK"] = os.environ["OMPI_COMM_WORLD_RANK"]
 
+        # Initialize distributed environment with SLURM
+        if "SLURM_PROCID" in os.environ.keys():
+            os.environ["RANK"] = os.environ["SLURM_PROCID"]
+
+        if "SLURM_NTASKS" in os.environ.keys():
+            os.environ["WORLD_SIZE"] = os.environ["SLURM_NTASKS"]
+
+        if "MASTER_ADDR" not in os.environ.keys():
+            os.environ["MASTER_ADDR"] = "127.0.0.1"
+
+        os.environ["MASTER_PORT"] = "1234"
         dist.init_process_group(backend='nccl')
         rank = dist.get_rank()
         size = dist.get_world_size()
@@ -738,7 +752,7 @@ def main():
 
     if graphname == "Cora":
         path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', graphname)
-        dataset = Planetoid(path, graphname, T.NormalizeFeatures())
+        dataset = Planetoid(path, graphname, transform=T.NormalizeFeatures())
         data = dataset[0]
         data = data.to(device)
         data.x.requires_grad = True

gcn_distr_2d.py

@@ -1095,6 +1095,9 @@ def twod_partition(rank, size, inputs, adj_matrix, data, features, classes, devi
     proc_row = proc_row_size(size)
     proc_col = proc_col_size(size)
 
+    inputs = inputs.to(torch.device("cpu"))
+    adj_matrix = adj_matrix.to(torch.device("cpu"))
+
     # n_per_proc = math.ceil(float(node_count) / proc_row)
     n_per_proc = node_count // proc_row
 
@@ -1401,7 +1404,7 @@ def main():
 
     # mid_layer = 16
     if graphname == 'Cora':
-        dataset = Planetoid(path, graphname, T.NormalizeFeatures())
+        dataset = Planetoid(path, graphname, transform=T.NormalizeFeatures())
         data = dataset[0]
         num_features = dataset.num_features
         num_classes = dataset.num_classes
@@ -1463,6 +1466,18 @@ def main():
 
     if "OMPI_COMM_WORLD_RANK" in os.environ.keys():
         os.environ["RANK"] = os.environ["OMPI_COMM_WORLD_RANK"]
+
+    # Initialize distributed environment with SLURM
+    if "SLURM_PROCID" in os.environ.keys():
+        os.environ["RANK"] = os.environ["SLURM_PROCID"]
+
+    if "SLURM_NTASKS" in os.environ.keys():
+        os.environ["WORLD_SIZE"] = os.environ["SLURM_NTASKS"]
+
+    if "MASTER_ADDR" not in os.environ.keys():
+        os.environ["MASTER_ADDR"] = "127.0.0.1"
+
+    os.environ["MASTER_PORT"] = "1234"
     dist.init_process_group(backend='nccl')
     # dist.init_process_group('gloo', init_method='env://')
     rank = dist.get_rank()