refractor the position of register.py

Author: JiaxuanYou

graphgym/config.py

@@ -4,8 +4,7 @@
 
 from graphgym.utils.io import makedirs_rm_exist
 
-from graphgym.contrib.config import *
-import graphgym.models.register as register
+import graphgym.register as register
 
 # Global config object
 cfg = CN()

graphgym/contrib/act/example.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 from graphgym.config import cfg
-from graphgym.models.register import register_act
+from graphgym.register import register_act
 
 
 class SWISH(nn.Module):

graphgym/contrib/config/example.py

@@ -1,6 +1,6 @@
 from yacs.config import CfgNode as CN
 
-from graphgym.models.register import register_config
+from graphgym.register import register_config
 
 
 def set_cfg_example(cfg):

graphgym/contrib/feature_augment/example.py

@@ -1,6 +1,6 @@
 import networkx as nx
 
-from graphgym.models.register import register_feature_augment
+from graphgym.register import register_feature_augment
 
 
 def example_node_augmentation_func(graph, **kwargs):

graphgym/contrib/feature_encoder/example.py

@@ -1,8 +1,6 @@
 import torch
-import torch.nn as nn
-from graphgym.config import cfg
-from graphgym.models.register import (register_node_encoder,
-                                      register_edge_encoder)
+from graphgym.register import (register_node_encoder,
+                               register_edge_encoder)
 
 from ogb.utils.features import get_bond_feature_dims
 

graphgym/contrib/head/example.py

@@ -1,6 +1,6 @@
 import torch.nn as nn
 
-from graphgym.models.register import register_head
+from graphgym.register import register_head
 
 
 class ExampleNodeHead(nn.Module):

graphgym/contrib/layer/attconv.py

@@ -8,7 +8,7 @@
 
 from torch_geometric.nn.inits import glorot, zeros
 from graphgym.config import cfg
-from graphgym.models.register import register_layer
+from graphgym.register import register_layer
 
 
 class GeneralAddAttConvLayer(MessagePassing):

graphgym/contrib/layer/example.py

@@ -5,7 +5,7 @@
 
 from torch_geometric.nn.inits import glorot, zeros
 from graphgym.config import cfg
-from graphgym.models.register import register_layer
+from graphgym.register import register_layer
 
 
 # Note: A registered GNN layer should take 'batch' as input

graphgym/contrib/layer/generalconv.py

@@ -1,14 +1,12 @@
 import torch
 import torch.nn as nn
 from torch.nn import Parameter
-import torch.nn.functional as F
 from torch_scatter import scatter_add
 from torch_geometric.nn.conv import MessagePassing
-from torch_geometric.utils import add_remaining_self_loops, softmax
+from torch_geometric.utils import add_remaining_self_loops
 
 from torch_geometric.nn.inits import glorot, zeros
 from graphgym.config import cfg
-from graphgym.models.register import register_layer
 
 
 class GeneralConvLayer(MessagePassing):

graphgym/contrib/layer/generalconv_ogb.py

@@ -8,9 +8,9 @@
 from torch_geometric.nn.inits import glorot, zeros
 from graphgym.config import cfg
 
-from graphgym.models.register import register_layer
+from graphgym.register import register_layer
 
-from ogb.utils.features import get_atom_feature_dims, get_bond_feature_dims
+from ogb.utils.features import get_bond_feature_dims
 
 full_bond_feature_dims = get_bond_feature_dims()
 

graphgym/contrib/layer/idconv.py

@@ -10,7 +10,7 @@
 from torch_geometric.nn.inits import glorot, zeros, reset
 
 from graphgym.config import cfg
-from graphgym.models.register import register_layer
+from graphgym.register import register_layer
 
 
 class GeneralIDConvLayer(MessagePassing):

graphgym/contrib/layer/sageinitconv.py

@@ -5,8 +5,8 @@
 from torch_geometric.nn.conv import MessagePassing
 from torch_geometric.utils import add_remaining_self_loops
 
-from torch_geometric.nn.inits import glorot, zeros, reset
-from graphgym.models.register import register_layer
+from torch_geometric.nn.inits import glorot, zeros
+from graphgym.register import register_layer
 
 
 class SAGEConvLayer(MessagePassing):

graphgym/contrib/loader/example.py

@@ -1,7 +1,7 @@
 from deepsnap.dataset import GraphDataset
 from torch_geometric.datasets import *
 
-from graphgym.models.register import register_loader
+from graphgym.register import register_loader
 
 
 def load_dataset_example(format, name, dataset_dir):

graphgym/contrib/loss/example.py

@@ -1,6 +1,6 @@
 import torch.nn as nn
 
-from graphgym.models.register import register_loss
+from graphgym.register import register_loss
 
 from graphgym.config import cfg
 

graphgym/contrib/network/example.py

@@ -5,7 +5,7 @@
 
 from graphgym.config import cfg
 from graphgym.models.head import head_dict
-from graphgym.models.register import register_network
+from graphgym.register import register_network
 
 
 class ExampleGNN(torch.nn.Module):

graphgym/contrib/optimizer/example.py

@@ -1,6 +1,6 @@
 import torch.optim as optim
 
-from graphgym.models.register import register_optimizer, register_scheduler
+from graphgym.register import register_optimizer, register_scheduler
 
 from graphgym.config import cfg
 

graphgym/contrib/pooling/example.py

@@ -1,8 +1,5 @@
-import torch
-import torch.nn as nn
 from torch_scatter import scatter
-from graphgym.config import cfg
-from graphgym.models.register import register_pooling
+from graphgym.register import register_pooling
 
 
 def global_example_pool(x, batch, size=None):

graphgym/contrib/stage/example.py

@@ -1,9 +1,7 @@
-import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch_scatter import scatter
 from graphgym.config import cfg
-from graphgym.models.register import register_stage
+from graphgym.register import register_stage
 
 import graphgym.models.gnn as gnn
 

graphgym/loader.py

@@ -1,4 +1,3 @@
-import os
 import networkx as nx
 import time
 import logging
@@ -15,14 +14,11 @@
 import graphgym.models.feature_augment as preprocess
 from graphgym.models.transform import (ego_nets, remove_node_feature,
                                        edge_nets, path_len)
-from graphgym.contrib.loader import *
-import graphgym.models.register as register
+import graphgym.register as register
 
 from ogb.graphproppred import PygGraphPropPredDataset
 from deepsnap.batch import Batch
 
-import pdb
-
 
 def load_pyg(name, dataset_dir):
     '''

graphgym/loss.py

@@ -2,8 +2,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from graphgym.contrib.loss import *
-import graphgym.models.register as register
+import graphgym.register as register
 from graphgym.config import cfg
 
 def compute_loss(pred, true):

graphgym/model_builder.py

@@ -3,8 +3,7 @@
 from graphgym.config import cfg
 from graphgym.models.gnn import GNN
 
-from graphgym.contrib.network import *
-import graphgym.models.register as register
+import graphgym.register as register
 
 network_dict = {
     'gnn': GNN,

graphgym/models/act.py

@@ -1,9 +1,6 @@
-import torch
 import torch.nn as nn
-import torch.nn.functional as F
 from graphgym.config import cfg
-from graphgym.contrib.act import *
-import graphgym.models.register as register
+import graphgym.register as register
 
 act_dict = {
     'relu': nn.ReLU(inplace=cfg.mem.inplace),

graphgym/models/feature_augment.py

@@ -9,7 +9,7 @@
 from graphgym.config import cfg
 from graphgym.contrib.transform.identity import compute_identity
 
-import graphgym.models.register as register
+import graphgym.register as register
 
 
 def _key(key, as_label=False):

graphgym/models/feature_encoder.py

@@ -1,8 +1,7 @@
 import torch
 from ogb.utils.features import get_atom_feature_dims, get_bond_feature_dims
 
-from graphgym.contrib.feature_encoder import *
-import graphgym.models.register as register
+import graphgym.register as register
 
 # Used for the OGB Encoders
 full_atom_feature_dims = get_atom_feature_dims()

graphgym/models/gnn.py

@@ -11,8 +11,7 @@
 from graphgym.init import init_weights
 from graphgym.models.feature_encoder import node_encoder_dict, edge_encoder_dict
 
-from graphgym.contrib.stage import *
-import graphgym.models.register as register
+import graphgym.register as register
 
 
 ########### Layer ############

graphgym/models/head.py

@@ -5,16 +5,12 @@
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
 from graphgym.config import cfg
-from graphgym.models.layer import layer_dict, MLP, GeneralLayer
+from graphgym.models.layer import MLP
 from graphgym.models.pooling import pooling_dict
 
-from graphgym.contrib.head import *
-import graphgym.models.register as register
-
-import pdb
+import graphgym.register as register
 
 
 ########### Head ############
@@ -61,8 +57,8 @@ def __init__(self, dim_in, dim_out):
                                      num_layers=cfg.gnn.layers_post_mp,
                                      bias=True)
             # requires parameter
-            self.decode_module = lambda v1, v2: torch.sigmoid(
-                self.layer_post_mp(torch.cat((v1, v2), dim=-1)))
+            self.decode_module = lambda v1, v2: \
+                self.layer_post_mp(torch.cat((v1, v2), dim=-1))
         else:
             if dim_out > 1:
                 raise ValueError(

graphgym/models/layer.py

@@ -8,8 +8,7 @@
 from graphgym.contrib.layer.generalconv import (GeneralConvLayer,
                                                 GeneralEdgeConvLayer)
 
-from graphgym.contrib.layer import *
-import graphgym.models.register as register
+import graphgym.register as register
 
 
 ## General classes

graphgym/models/pooling.py

@@ -2,8 +2,7 @@
 from torch_scatter import scatter
 from graphgym.config import cfg
 
-from graphgym.contrib.pooling import *
-import graphgym.models.register as register
+import graphgym.register as register
 
 
 # Pooling options (pool nodes into graph representations)

graphgym/optimizer.py

@@ -2,8 +2,7 @@
 
 import torch.optim as optim
 
-from graphgym.contrib.optimizer import *
-import graphgym.models.register as register
+import graphgym.register as register
 
 
 def create_optimizer(params):