torch_geometric.graphgym
Workflow and Register Modules
Loads the model checkpoint at a given epoch. |
|
Saves the model checkpoint at a given epoch. |
|
Removes the model checkpoint at a given epoch. |
|
Removes all but the last model checkpoint. |
|
Parses the command line arguments. |
|
|
|
This function sets the default config value. |
|
Load configurations from file system and command line |
|
Dumps the config to the output directory specified in |
|
Create the directory for each random seed experiment run |
|
Create the directory for full experiment run |
|
Extract filename from file name path |
|
Performs weight initialization |
|
Create data loader object |
|
Set up printing options |
|
Create logger for the experiment. |
|
Compute loss and prediction score |
|
Create model for graph machine learning. |
|
Creates a config-driven optimizer. |
|
Creates a config-driven learning rate scheduler. |
|
|
|
Base function for registering a module in GraphGym. |
|
Registers an activation function in GraphGym. |
|
Registers a node feature encoder in GraphGym. |
|
Registers an edge feature encoder in GraphGym. |
|
Registers a customized GNN stage in GraphGym. |
|
Registers a GNN prediction head in GraphGym. |
|
Registers a GNN layer in GraphGym. |
|
Registers a GNN global pooling/readout layer in GraphGym. |
|
Registers a GNN model in GraphGym. |
|
Registers a configuration group in GraphGym. |
|
Registers a dataset in GraphGym. |
|
Registers a data loader in GraphGym. |
|
Registers an optimizer in GraphGym. |
|
Registers a learning rate scheduler in GraphGym. |
|
Registers a loss function in GraphGym. |
|
Registers a training function in GraphGym. |
|
Register a metric function in GraphGym. |
- load_ckpt(model: Module, optimizer: Optional[Optimizer] = None, scheduler: Optional[Any] = None, epoch: int = -1) int [source]
Loads the model checkpoint at a given epoch.
- save_ckpt(model: Module, optimizer: Optional[Optimizer] = None, scheduler: Optional[Any] = None, epoch: int = 0)[source]
Saves the model checkpoint at a given epoch.
- set_cfg(cfg)[source]
This function sets the default config value. 1) Note that for an experiment, only part of the arguments will be used The remaining unused arguments won’t affect anything. So feel free to register any argument in graphgym.contrib.config 2) We support at most two levels of configs, e.g., cfg.dataset.name
- Returns
configuration use by the experiment.
- load_cfg(cfg, args)[source]
Load configurations from file system and command line
- Parameters
cfg (CfgNode) – Configuration node
args (ArgumentParser) – Command argument parser
- dump_cfg(cfg)[source]
Dumps the config to the output directory specified in
cfg.out_dir
- Parameters
cfg (CfgNode) – Configuration node
- set_run_dir(out_dir)[source]
Create the directory for each random seed experiment run
- Parameters
out_dir (string) – Directory for output, specified in
cfg.out_dir
fname (string) – Filename for the yaml format configuration file
- set_out_dir(out_dir, fname)[source]
Create the directory for full experiment run
- Parameters
out_dir (string) – Directory for output, specified in
cfg.out_dir
fname (string) – Filename for the yaml format configuration file
- get_fname(fname)[source]
Extract filename from file name path
- Parameters
fname (string) – Filename for the yaml format configuration file
- compute_loss(pred, true)[source]
Compute loss and prediction score
- Parameters
pred (torch.tensor) – Unnormalized prediction
true (torch.tensor) – Grou
Returns: Loss, normalized prediction score
- create_model(to_device=True, dim_in=None, dim_out=None) GraphGymModule [source]
Create model for graph machine learning.
- create_optimizer(params: Iterator[Parameter], cfg: Any) Any [source]
Creates a config-driven optimizer.
- create_scheduler(optimizer: Optimizer, cfg: Any) Any [source]
Creates a config-driven learning rate scheduler.
- register_base(mapping: Dict[str, Any], key: str, module: Optional[Any] = None) Union[None, Callable] [source]
Base function for registering a module in GraphGym.
- register_act(key: str, module: Optional[Any] = None)[source]
Registers an activation function in GraphGym.
- register_node_encoder(key: str, module: Optional[Any] = None)[source]
Registers a node feature encoder in GraphGym.
- register_edge_encoder(key: str, module: Optional[Any] = None)[source]
Registers an edge feature encoder in GraphGym.
- register_stage(key: str, module: Optional[Any] = None)[source]
Registers a customized GNN stage in GraphGym.
- register_head(key: str, module: Optional[Any] = None)[source]
Registers a GNN prediction head in GraphGym.
- register_pooling(key: str, module: Optional[Any] = None)[source]
Registers a GNN global pooling/readout layer in GraphGym.
- register_network(key: str, module: Optional[Any] = None)[source]
Registers a GNN model in GraphGym.
- register_config(key: str, module: Optional[Any] = None)[source]
Registers a configuration group in GraphGym.
- register_loader(key: str, module: Optional[Any] = None)[source]
Registers a data loader in GraphGym.
- register_optimizer(key: str, module: Optional[Any] = None)[source]
Registers an optimizer in GraphGym.
- register_scheduler(key: str, module: Optional[Any] = None)[source]
Registers a learning rate scheduler in GraphGym.
- register_loss(key: str, module: Optional[Any] = None)[source]
Registers a loss function in GraphGym.
Model Modules
Provides an encoder for integer node features. |
|
The atom Encoder used in OGB molecule dataset. |
|
The bond Encoder used in OGB molecule dataset. |
|
Wrapper for a GNN layer |
|
Wrapper for NN layer before GNN message passing |
|
Simple Stage that stack GNN layers |
|
Encoding node and edge features |
|
General GNN model: encoder + stage + head |
|
GNN prediction head for node prediction tasks. |
|
GNN prediction head for edge/link prediction tasks. |
|
GNN prediction head for graph prediction tasks. |
|
General wrapper for layers |
|
General wrapper for a stack of multiple layers |
|
Basic Linear layer. |
|
BatchNorm for node feature. |
|
BatchNorm for edge feature. |
|
Basic MLP model. |
|
Graph Convolutional Network (GCN) layer |
|
GraphSAGE Conv layer |
|
Graph Attention Network (GAT) layer |
|
Graph Isomorphism Network (GIN) layer |
|
SplineCNN layer |
|
A general GNN layer |
|
A general GNN layer that supports edge features as well |
|
A general GNN layer that supports edge features and edge sampling |
|
Returns batch-wise graph-level-outputs by adding node features across the node dimension, so that for a single graph \(\mathcal{G}_i\) its output is computed by |
|
Returns batch-wise graph-level-outputs by averaging node features across the node dimension, so that for a single graph \(\mathcal{G}_i\) its output is computed by |
|
Returns batch-wise graph-level-outputs by taking the channel-wise maximum across the node dimension, so that for a single graph \(\mathcal{G}_i\) its output is computed by |
- class IntegerFeatureEncoder(emb_dim, num_classes=None)[source]
Provides an encoder for integer node features.
- class AtomEncoder(emb_dim, num_classes=None)[source]
The atom Encoder used in OGB molecule dataset.
- Parameters
emb_dim (int) – Output embedding dimension
num_classes – None
- class BondEncoder(emb_dim)[source]
The bond Encoder used in OGB molecule dataset.
- Parameters
emb_dim (int) – Output edge embedding dimension
- class FeatureEncoder(dim_in)[source]
Encoding node and edge features
- Parameters
dim_in (int) – Input feature dimension
- class GNNGraphHead(dim_in, dim_out)[source]
GNN prediction head for graph prediction tasks. The optional post_mp layer (specified by cfg.gnn.post_mp) is used to transform the pooled embedding using an MLP.
- class GeneralLayer(name, layer_config: LayerConfig, **kwargs)[source]
General wrapper for layers
- Parameters
name (string) – Name of the layer in registered
layer_dict
dim_in (int) – Input dimension
dim_out (int) – Output dimension
has_act (bool) – Whether has activation after the layer
has_bn (bool) – Whether has BatchNorm in the layer
has_l2norm (bool) – Wheter has L2 normalization after the layer
**kwargs (optional) – Additional args
- class GeneralMultiLayer(name, layer_config: LayerConfig, **kwargs)[source]
General wrapper for a stack of multiple layers
- Parameters
name (string) – Name of the layer in registered
layer_dict
num_layers (int) – Number of layers in the stack
dim_in (int) – Input dimension
dim_out (int) – Output dimension
dim_inner (int) – The dimension for the inner layers
final_act (bool) – Whether has activation after the layer stack
**kwargs (optional) – Additional args
- class BatchNorm1dNode(layer_config: LayerConfig)[source]
BatchNorm for node feature.
- Parameters
dim_in (int) – Input dimension
- class BatchNorm1dEdge(layer_config: LayerConfig)[source]
BatchNorm for edge feature.
- Parameters
dim_in (int) – Input dimension
- class MLP(layer_config: LayerConfig, **kwargs)[source]
Basic MLP model. Here 1-layer MLP is equivalent to a Liner layer.
- class GeneralEdgeConv(layer_config: LayerConfig, **kwargs)[source]
A general GNN layer that supports edge features as well
- class GeneralSampleEdgeConv(layer_config: LayerConfig, **kwargs)[source]
A general GNN layer that supports edge features and edge sampling
- global_add_pool(x: Tensor, batch: Optional[Tensor], size: Optional[int] = None) Tensor [source]
Returns batch-wise graph-level-outputs by adding node features across the node dimension, so that for a single graph \(\mathcal{G}_i\) its output is computed by
\[\mathbf{r}_i = \sum_{n=1}^{N_i} \mathbf{x}_n.\]Functional method of the
SumAggregation
module.- Parameters
x (Tensor) – Node feature matrix \(\mathbf{X} \in \mathbb{R}^{(N_1 + \ldots + N_B) \times F}\).
batch (LongTensor, optional) – Batch vector \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N\), which assigns each node to a specific example.
size (int, optional) – Batch-size \(B\). Automatically calculated if not given. (default:
None
)
- global_mean_pool(x: Tensor, batch: Optional[Tensor], size: Optional[int] = None) Tensor [source]
Returns batch-wise graph-level-outputs by averaging node features across the node dimension, so that for a single graph \(\mathcal{G}_i\) its output is computed by
\[\mathbf{r}_i = \frac{1}{N_i} \sum_{n=1}^{N_i} \mathbf{x}_n.\]Functional method of the
MeanAggregation
module.- Parameters
x (Tensor) – Node feature matrix \(\mathbf{X} \in \mathbb{R}^{(N_1 + \ldots + N_B) \times F}\).
batch (LongTensor, optional) – Batch vector \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N\), which assigns each node to a specific example.
size (int, optional) – Batch-size \(B\). Automatically calculated if not given. (default:
None
)
- global_max_pool(x: Tensor, batch: Optional[Tensor], size: Optional[int] = None) Tensor [source]
Returns batch-wise graph-level-outputs by taking the channel-wise maximum across the node dimension, so that for a single graph \(\mathcal{G}_i\) its output is computed by
\[\mathbf{r}_i = \mathrm{max}_{n=1}^{N_i} \, \mathbf{x}_n.\]Functional method of the
MaxAggregation
module.- Parameters
x (Tensor) – Node feature matrix \(\mathbf{X} \in \mathbb{R}^{(N_1 + \ldots + N_B) \times F}\).
batch (LongTensor, optional) – Batch vector \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N\), which assigns each node to a specific example.
size (int, optional) – Batch-size \(B\). Automatically calculated if not given. (default:
None
)
Utility Modules
Aggregate over different random seeds of a single experiment |
|
Aggregate across results from multiple experiments via grid search |
|
Computes the number of parameters. |
|
Match the computational budget of a given baseline model. |
|
Get the current GPU memory usage. |
|
Auto select device for the current experiment. |
|
Determines if the model should be evaluated at the current epoch. |
|
Determines if the model should be evaluated at the current epoch. |
|
Dump a Python dictionary to JSON file |
|
Dump a list of Python dictionaries to JSON file |
|
Add a dictionary of statistics to a Tensorboard writer |
|
Make a directory, remove any existing data. |
|
Default context manager that does nothing |
- agg_runs(dir, metric_best='auto')[source]
Aggregate over different random seeds of a single experiment
- agg_batch(dir, metric_best='auto')[source]
Aggregate across results from multiple experiments via grid search
- params_count(model)[source]
Computes the number of parameters.
- Parameters
model (nn.Module) – PyTorch model
- match_baseline_cfg(cfg_dict, cfg_dict_baseline, verbose=True)[source]
Match the computational budget of a given baseline model. THe current configuration dictionary will be modifed and returned.
- dict_list_to_json(dict_list, fname)[source]
Dump a list of Python dictionaries to JSON file
- Parameters
dict_list (list of dict) – List of Python dictionaries
fname (str) – Output file name