Source code for torch_geometric.graphgym.register

def register(key, module, module_dict):
    """
    Base function for registering a customized module to a module dictionary

    Args:
        key (string): Name of the module
        module: PyTorch module
        module_dict (dict): Python dictionary,
            hosting all the registered modules

    """
    if key in module_dict:
        raise KeyError('Key {} is already pre-defined.'.format(key))
    else:
        module_dict[key] = module


act_dict = {}


def register_act(key, module):
    """
    Register a customized activation function.
    After registeration, the module can be directly called by GraphGym.

    Args:
        key (string): Name of the module
        module: PyTorch module

    """
    register(key, module, act_dict)


node_encoder_dict = {}


def register_node_encoder(key, module):
    """
    Register a customized node feature encoder.
    After registeration, the module can be directly called by GraphGym.

    Args:
        key (string): Name of the module
        module: PyTorch module

    """
    register(key, module, node_encoder_dict)


edge_encoder_dict = {}


def register_edge_encoder(key, module):
    """
    Register a customized edge feature encoder.
    After registeration, the module can be directly called by GraphGym.

    Args:
        key (string): Name of the module
        module: PyTorch module

    """
    register(key, module, edge_encoder_dict)


stage_dict = {}


def register_stage(key, module):
    """
    Register a customized GNN stage (consists of multiple layers).
    After registeration, the module can be directly called by GraphGym.

    Args:
        key (string): Name of the module
        module: PyTorch module

    """
    register(key, module, stage_dict)


head_dict = {}


def register_head(key, module):
    """
    Register a customized GNN prediction head.
    After registeration, the module can be directly called by GraphGym.

    Args:
        key (string): Name of the module
        module: PyTorch module

    """
    register(key, module, head_dict)


layer_dict = {}


def register_layer(key, module):
    """
    Register a customized GNN layer.
    After registeration, the module can be directly called by GraphGym.

    Args:
        key (string): Name of the module
        module: PyTorch module

    """
    register(key, module, layer_dict)


pooling_dict = {}


def register_pooling(key, module):
    """
    Register a customized GNN pooling layer (for graph classification).
    After registeration, the module can be directly called by GraphGym.

    Args:
        key (string): Name of the module
        module: PyTorch module

    """
    register(key, module, pooling_dict)


network_dict = {}


def register_network(key, module):
    """
    Register a customized GNN model.
    After registeration, the module can be directly called by GraphGym.

    Args:
        key (string): Name of the module
        module: PyTorch module

    """
    register(key, module, network_dict)


config_dict = {}


def register_config(key, module):
    """
    Register a customized configuration group.
    After registeration, the module can be directly called by GraphGym.

    Args:
        key (string): Name of the module
        module: PyTorch module

    """
    register(key, module, config_dict)


loader_dict = {}


def register_loader(key, module):
    """
    Register a customized PyG data loader.
    After registeration, the module can be directly called by GraphGym.

    Args:
        key (string): Name of the module
        module: PyTorch module

    """
    register(key, module, loader_dict)


optimizer_dict = {}


def register_optimizer(key, module):
    """
    Register a customized optimizer.
    After registeration, the module can be directly called by GraphGym.

    Args:
        key (string): Name of the module
        module: PyTorch module

    """
    register(key, module, optimizer_dict)


scheduler_dict = {}


def register_scheduler(key, module):
    """
    Register a customized learning rate scheduler.
    After registeration, the module can be directly called by GraphGym.

    Args:
        key (string): Name of the module
        module: PyTorch module

    """
    register(key, module, scheduler_dict)


loss_dict = {}


def register_loss(key, module):
    """
    Register a customized loss function.
    After registeration, the module can be directly called by GraphGym.

    Args:
        key (string): Name of the module
        module: PyTorch module

    """
    register(key, module, loss_dict)


train_dict = {}


def register_train(key, module):
    """
    Register a customized training function.
    After registeration, the module can be directly called by GraphGym.

    Args:
        key (string): Name of the module
        module: PyTorch module

    """
    register(key, module, train_dict)