charlie

root_gnn_dgl/models/GCN.py

@@ -1154,6 +1154,7 @@ class Attention(nn.Module):
         self.n_proc_steps = n_proc_steps
         self.layers = nn.ModuleList()
         self.has_global = sample_global.shape[1] != 0
+        self.hid_size = hid_size
         gl_size = sample_global.shape[1] if self.has_global else 1
 
         #encoder
@@ -1196,7 +1197,7 @@ class Attention(nn.Module):
                 batch_num_nodes.append(non_padded_count)
                 start_idx = end_idx
             batch_num_nodes = torch.tensor(batch_num_nodes, device = g.ndata['features'].device)
-            sum_weights = batch_num_nodes[:, None].repeat(1, 64)
+            sum_weights = batch_num_nodes[:, None].repeat(1, self.hid_size)
             global_feats = batch_num_nodes[:, None].to(torch.float)
 
         h_global = self.global_encoder(global_feats)
@@ -1364,6 +1365,7 @@ class Transferred_Learning_Attention(nn.Module):
         self.n_proc_steps = n_proc_steps
         self.layers = nn.ModuleList()
         self.has_global = sample_global.shape[1] != 0
+        self.hid_size = hid_size
         gl_size = sample_global.shape[1] if self.has_global else 1
 
         self.learning_rate = learning_rate
@@ -1440,7 +1442,7 @@ class Transferred_Learning_Attention(nn.Module):
                 batch_num_nodes.append(non_padded_count)
                 start_idx = end_idx
             batch_num_nodes = torch.tensor(batch_num_nodes, device = g.ndata['features'].device)
-            sum_weights = batch_num_nodes[:, None].repeat(1, 64)
+            sum_weights = batch_num_nodes[:, None].repeat(1, self.hid_size)
             global_feats = batch_num_nodes[:, None].to(torch.float)
 
         h_global = self.TL_global_encoder(global_feats)
@@ -1856,6 +1858,7 @@ class Clustering(nn.Module):
         self.n_layers = n_layers
         self.n_proc_steps = n_proc_steps
         self.layers = nn.ModuleList()
+        self.hid_size = hid_size
         if (len(sample_global) == 0):
             self.has_global = False
         else:
@@ -1899,7 +1902,7 @@ class Clustering(nn.Module):
                 batch_num_nodes.append(non_padded_count)
                 start_idx = end_idx
             batch_num_nodes = torch.tensor(batch_num_nodes, device = g.ndata[features].device)
-            sum_weights = batch_num_nodes[:, None].repeat(1, 64)
+            sum_weights = batch_num_nodes[:, None].repeat(1, self.hid_size)
             global_feats = batch_num_nodes[:, None].to(torch.float)
 
         h_global = self.global_encoder(global_feats)

root_gnn_dgl/root_gnn_base/batched_dataset.py

@@ -16,7 +16,7 @@ def GetBatchedLoader(dataset, batch_size, mask_fn = None, drop_last=True, **kwar
 
 #Dataset which contains prebatched shuffled graphs. Cannot be saved to disk, else batching info is lost.
 class PreBatchedDataset(DGLDataset):
-    def __init__(self, start_dataset, batch_size, mask_fn = None, drop_last=True, save_to_disk = True, suffix = '', chunks = 1, chunkno = -1, shuffle = True, padding_mode = 'NONE', **kwargs):
+    def __init__(self, start_dataset, batch_size, mask_fn = None, drop_last=True, save_to_disk = True, suffix = '', chunks = 1, chunkno = -1, shuffle = True, padding_mode = 'NONE', hidden_size=64, **kwargs):
         print(f'Unused kwargs: {kwargs}')
         self.start_dataset = start_dataset
         self.start_dataset.load()
@@ -34,6 +34,7 @@ class PreBatchedDataset(DGLDataset):
         self.suffix = suffix
         self.current_chunk = None
         self.current_chunk_idx = -1
+        self.hid_size = hidden_size
         super().__init__(name = start_dataset.name + '_prebatched_padded', save_dir=start_dataset.save_dir)
 
     def process(self):
@@ -86,7 +87,7 @@ class PreBatchedDataset(DGLDataset):
             for i in range(len(self.graphs)):
                 unbatched_g = dgl.unbatch(self.graphs[i])
                 max_num_nodes = max(g.number_of_nodes() for g in unbatched_g)
-                self.graphs[i] = utils.pad_batch_num_nodes(self.graphs[i], max_num_nodes)
+                self.graphs[i] = utils.pad_batch_num_nodes(self.graphs[i], max_num_nodes, hid_size=self.hid_size)
                 self.batch_num_nodes.append(self.graphs[i].batch_num_nodes())
                 self.batch_num_edges.append(self.graphs[i].batch_num_edges())
         else:

root_gnn_dgl/root_gnn_base/dataset.py

@@ -162,6 +162,7 @@ class RootDataset(DGLDataset):
         self.global_chunks = []
         chunk_id = -1
         for chunk in chunks:
+            print('Processing chunk {}/{}'.format(chunk_id + 1, len(chunks)), flush=True)
             chunk_id += 1
             graphs = []
             labels = []
@@ -198,6 +199,12 @@ class RootDataset(DGLDataset):
             labels = torch.tensor(labels)
             tracking = torch.stack(tracking)
             globals = torch.stack(globals)
+
+            self.graph_chunks.append(graphs)
+            self.label_chunks.append(labels)
+            self.tracking_chunks.append(tracking)
+            self.global_chunks.append(globals)
+            self.counts.append(len(graphs))
             
             if (self.chunks > 1):
                 self.save_chunk(chunk_id, graphs, labels, tracking, globals)
@@ -223,16 +230,12 @@ class RootDataset(DGLDataset):
             return
         graph_path = os.path.join(self.save_dir, self.name + '.bin')
         if self.chunks == 1:
-            # print(len(self.graphs))
-            # print(len(self.labels))
-            # print(len(self.tracking))
-            # print(len(self.globals))
             print(f'Saving dataset to {os.path.join(self.save_dir, self.name + ".bin")}')
             dgl.save_graphs(str(graph_path), self.graphs, {'labels': torch.tensor(self.labels), 'tracking': torch.tensor(self.tracking), 'global': torch.tensor(self.global_features)})
         else:
-            print(len(self.graph_chunks))
             for i in range(len(self.process_chunks)):
                 print(f'Saving dataset to {os.path.join(self.save_dir, self.name + f"_{self.process_chunks[i]}.bin")}')
+
                 dgl.save_graphs(str(graph_path).replace('.bin', f'_{self.process_chunks[i]}.bin'), self.graph_chunks[i], {'labels': self.label_chunks[i], 'tracking': self.tracking_chunks[i], 'global': self.global_chunks[i]})
 
     def save_chunk(self, chunk_id, graphs, labels, tracking, globals):

root_gnn_dgl/root_gnn_base/utils.py

@@ -92,7 +92,7 @@ def pad_batch(batch, edges = 104000, nodes = 16000):
     return make_padding_graph(batch, pad_nodes, pad_edges)
 
 def pad_batch_num_nodes(batch, max_num_nodes, hid_size = 64):
-    print(f"Padding each graph to have {max_num_nodes} nodes")
+    print(f"Padding each graph to have {max_num_nodes} nodes. Using hidden size {hid_size}.")
 
     unbatched = dgl.unbatch(batch)
     for g in unbatched:

root_gnn_dgl/scripts/prep_data.py

@@ -33,12 +33,12 @@ def main():
         fold_conf = dset_config["folding"]
         print(f"shuffle_chunks = {shuffle_chunks}, args.chunk = {args.chunk}, padding_mode = {padding_mode}")
         if dset_config["class"] == "LazyMultiLabelDataset":
-            LazyPreBatchedDataset(start_dataset = dset, batch_size = batch_size, mask_fn = utils.fold_selection(fold_conf, "train"), suffix = utils.fold_selection_name(fold_conf, "train"), chunks = shuffle_chunks, chunkno = args.chunk, padding_mode = padding_mode, drop_last=args.drop_last)
-            LazyPreBatchedDataset(start_dataset = dset, batch_size = batch_size, mask_fn = utils.fold_selection(fold_conf, "test"),  suffix = utils.fold_selection_name(fold_conf, 'test'), chunks = shuffle_chunks, chunkno = args.chunk, padding_mode = padding_mode, drop_last=args.drop_last)
+            LazyPreBatchedDataset(start_dataset = dset, batch_size = batch_size, mask_fn = utils.fold_selection(fold_conf, "train"), suffix = utils.fold_selection_name(fold_conf, "train"), chunks = shuffle_chunks, chunkno = args.chunk, padding_mode = padding_mode, drop_last=args.drop_last, hidden_size=config['Model']['args']['hid_size'] )
+            LazyPreBatchedDataset(start_dataset = dset, batch_size = batch_size, mask_fn = utils.fold_selection(fold_conf, "test"),  suffix = utils.fold_selection_name(fold_conf, 'test'), chunks = shuffle_chunks, chunkno = args.chunk, padding_mode = padding_mode, drop_last=args.drop_last, hidden_size=config['Model']['args']['hid_size'])
 
         else:
-            PreBatchedDataset(dset, batch_size, utils.fold_selection(fold_conf, "train"), suffix = utils.fold_selection_name(fold_conf, "train"), chunks = shuffle_chunks, chunkno = args.chunk, padding_mode = padding_mode, drop_last=args.drop_last)
-            PreBatchedDataset(dset, batch_size, utils.fold_selection(fold_conf, "test"),  suffix = utils.fold_selection_name(fold_conf, 'test'), chunks = shuffle_chunks, chunkno = args.chunk, padding_mode = padding_mode, drop_last=args.drop_last)
+            PreBatchedDataset(dset, batch_size, utils.fold_selection(fold_conf, "train"), suffix = utils.fold_selection_name(fold_conf, "train"), chunks = shuffle_chunks, chunkno = args.chunk, padding_mode = padding_mode, drop_last=args.drop_last,hidden_size=config['Model']['args']['hid_size'])
+            PreBatchedDataset(dset, batch_size, utils.fold_selection(fold_conf, "test"),  suffix = utils.fold_selection_name(fold_conf, 'test'), chunks = shuffle_chunks, chunkno = args.chunk, padding_mode = padding_mode, drop_last=args.drop_last,hidden_size=config['Model']['args']['hid_size'] )
 
 if __name__ == "__main__":
-    main()
+    main()

root_gnn_dgl/scripts/training_script.py

@@ -707,7 +707,7 @@ def main(rank=0, args=None, world_size=1, port=24500, seed=12345):
         mask_fn = utils.fold_selection(fold_conf, "train")
         if args.preshuffle:
             # ldr = ldr_type(start_dataset=dset, batch_size=batch_size, mask_fn=mask_fn, suffix = utils.fold_selection_name(fold_conf, 'train'), chunks = shuffle_chunks, padding_mode = padding_mode, use_ddp = args.multigpu, rank=rank, world_size=world_size)
-            ldr = ldr_type(start_dataset=dset, batch_size=batch_size, mask_fn=mask_fn, suffix = utils.fold_selection_name(fold_conf, 'train'), chunks = shuffle_chunks, padding_mode = padding_mode)
+            ldr = ldr_type(start_dataset=dset, batch_size=batch_size, mask_fn=mask_fn, suffix = utils.fold_selection_name(fold_conf, 'train'), chunks = shuffle_chunks, padding_mode = padding_mode, hidden_size = config["Model"]["args"]["hid_size"])
             gsamp, _, _, global_samp = ldr[0]
             sampler = None
             
@@ -724,7 +724,7 @@ def main(rank=0, args=None, world_size=1, port=24500, seed=12345):
                 sampler = DistributedSampler(ldr, num_replicas=world_size, rank=pargs.global_rank, shuffle=False, drop_last=True)
             train_loaders.append(torch.utils.data.DataLoader(ldr, batch_size = None, num_workers = 0, sampler = sampler))
             sampler = None
-            ldr = ldr_type(start_dataset=dset, batch_size=batch_size, mask_fn=mask_fn, suffix = utils.fold_selection_name(fold_conf, 'test'), chunks = shuffle_chunks, padding_mode = padding_mode)
+            ldr = ldr_type(start_dataset=dset, batch_size=batch_size, mask_fn=mask_fn, suffix = utils.fold_selection_name(fold_conf, 'test'), chunks = shuffle_chunks, padding_mode = padding_mode, hidden_size= config['Model']['args']['hid_size'])
             if (args.multigpu):
                 sampler = DistributedSampler(ldr, num_replicas=world_size, rank=rank, shuffle=False, drop_last=True)
                 # num_batches = len(ldr)
@@ -737,7 +737,7 @@ def main(rank=0, args=None, world_size=1, port=24500, seed=12345):
             test_loaders.append(torch.utils.data.DataLoader(ldr, batch_size = None, num_workers = 0, sampler=sampler))
 
             if "validation" in fold_conf:
-                val_loaders.append(torch.utils.data.DataLoader((ldr_type(start_dataset=dset, batch_size=batch_size, mask_fn=utils.fold_selection(fold_conf, "validation"), suffix = utils.fold_selection_name(fold_conf, 'validation'), chunks = shuffle_chunks, padding_mode = padding_mode, rank=rank, world_size=1)), batch_size = None, num_workers = 0, sampler = sampler))
+                val_loaders.append(torch.utils.data.DataLoader((ldr_type(start_dataset=dset, batch_size=batch_size, mask_fn=utils.fold_selection(fold_conf, "validation"), suffix = utils.fold_selection_name(fold_conf, 'validation'), chunks = shuffle_chunks, hidden_size=config['Model']['args']['hid_size'],  padding_mode = padding_mode, rank=rank, world_size=1)), batch_size = None, num_workers = 0, sampler = sampler))
             else:
                 print("No validation set for dataset ", dset_conf)
         else:
@@ -753,6 +753,8 @@ def main(rank=0, args=None, world_size=1, port=24500, seed=12345):
     print("Load time: {:.4f} s".format(load_end - load_start))
 
     model = utils.buildFromConfig(config["Model"], {'sample_graph': gsamp, 'sample_global': global_samp, 'seed': seed}).to(device)
+    pytorch_total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f"Number of trainable parameters = {pytorch_total_params}")
     if not args.nocompile:
         model = torch.compile(model)
     if args.multigpu: