kgcnn-torch

kgcnn-torch is a native PyTorch and PyTorch Geometric implementation of graph neural networks for molecular and materials-science prediction. It provides reusable layers, models, datasets, and training utilities using PyG's Data and Batch interfaces.

This is an independent PyTorch/PyG port inspired by the model definitions and data tooling in aimat-lab/kgcnn. It is not an official release of the original project.

• Native torch.nn.Module implementations designed for PyG data objects.
• Molecular, crystal, node-classification, and energy-and-force workflows.
• Periodic-boundary-condition support for crystal models.
• An EnergyForceModel wrapper that derives forces from predicted energies.
• Dataset loaders, preprocessing utilities, cross-validation, early stopping, checkpoints, learning-rate scheduling, metrics, and label scaling.
• Example notebooks and JSON-based training configurations.

The repository includes implementations of GCN, GAT, GATv2, GIN, rGIN, GraphSAGE, RGCN, SchNet, DimeNet++, DMPNN, CMPNN, DGIN, NMPN, INorp, PAiNN, EGNN, HamNet, AttentiveFP, MAT, MEGAN, MoGAT, CGCNN, MEGNet, HDNNP2nd, MXMNet, GNNFiLM, and GNNExplainer utilities. Some architectures also provide crystal-aware variants.

All layers live in kgcnn_torch/layers and use the PyG edge convention: edge_index[0] is the source node and edge_index[1] is the target node.

layers/geom.py contains the geometric building blocks used by molecular and crystal models:

Aggregation supports common reduction modes such as sum, mean, max, and min. The attentive pooling implementation includes the iterative GRU-and-attention refinement used by AttentiveFP.

MLP supports per-layer widths and activations, optional dropout, and batch, layer, graph, group, or unit normalization. Graph-aware normalization is provided by GraphBatchNorm, GraphLayerNorm, and GraphNormalization.

Other useful components include gather_nodes_outgoing, gather_nodes_ingoing, GRUUpdate, ResidualLayer, StandardLabelScaler, and ExtensiveMolecularLabelScaler. Low-level scatter reductions and activations are available under kgcnn_torch/ops, including scatter_reduce_sum, scatter_reduce_mean, scatter_reduce_max, scatter_reduce_min, and scatter_reduce_softmax.

The training module provides a framework-independent PyTorch training loop in kgcnn_torch/training:

from kgcnn_torch.training.trainer import fit

history = fit(
    model=model,
    train_loader=train_loader,
    val_loader=val_loader,
    optimizer=optimizer,
    loss_fn=loss_fn,
    epochs=500,
    metrics={"mae": mae, "rmse": rmse},
    scheduler=scheduler,
    callbacks=[early_stopping, checkpoint],
    device="cuda",
    scaler=scaler,
)

EarlyStoppingCallback, ModelCheckpointCallback, and LearningRateLoggingCallback provide the standard training lifecycle hooks. The scheduler factory supports linear warm-up, warm-up plus exponential or cosine decay, polynomial decay, linear decay, ReduceLROnPlateau, StepLR, ExponentialLR, and CosineAnnealingLR.

• EnergyForceLoss combines energy and force objectives.
• ForceMeanAbsoluteError and DisjointForceMeanAbsoluteError evaluate forces while handling molecule-wise normalization or disjoint graphs.
• BinaryCrossentropyNoNaN, BinaryAccuracyNoNaN, BalancedBinaryAccuracyNoNaN, and AUCNoNaN handle datasets with missing labels.
• Regression metrics include MAE, MSE, RMSE, and scaled MAE/RMSE variants.
• StandardLabelScaler applies standard scaling; ExtensiveMolecularLabelScaler learns element reference energies with ridge regression for extensive targets.

JSON configurations are loaded by HyperParameter from kgcnn_torch/training/hyper.py. A file can store multiple model configurations keyed by model name:

{
  "SchNet": {
    "model": {
      "config": {
        "num_features": 128,
        "num_filters": 128,
        "num_interactions": 6,
        "cutoff": 10.0,
        "num_gaussians": 50,
        "output_dim": 1
      }
    },
    "training": {
      "fit": {"epochs": 500, "batch_size": 32},
      "compile": {"optimizer": {"class_name": "Adam", "config": {"lr": 0.0005}}, "loss": "mae"},
      "cross_validation": {"n_splits": 5, "shuffle": true}
    }
  }
}

The training_scripts/hyper directory contains ready-to-run configurations for ESOL, FreeSolv, Lipophilicity, QM7/QM9, MUTAG, Mutagenicity, PROTEINS, ClinTox, SIDER, Tox21, Cora, MD17, ISO17, MatBench, and Materials Project tasks.

For GNNFilm and RGCN, output_final_activation defaults to "linear" to work naturally with logits-based losses such as BCEWithLogitsLoss and CrossEntropyLoss. Set it explicitly to "softmax" when reproducing a configuration that expects probability outputs.

PyG Data is the standard graph representation. The dataset pipeline can convert MemoryGraphList-style kgcnn data into PyG objects.

Edge-order note: the original kgcnn convention is [target, source], while PyG uses [source, target]. to_pyg_list() performs this conversion.

The graph, molecule, crystal, and io modules cover generic graph preprocessing, RDKit/Open Babel molecule processing, periodic crystal graph construction, and data I/O respectively.

Python 3.9 or later is required. Install a PyTorch build appropriate for your CPU or CUDA environment first, then install PyG following its official installation guide.

git clone https://github.com/DinyangLyu/kgcnn-torch.git
cd kgcnn-torch
pip install -e .

Optional dependency groups are available for common workflows:

pip install -e ".[molecule]"  # RDKit, Open Babel, ASE
pip install -e ".[crystal]"   # pymatgen, pyxtal
pip install -e ".[vis]"       # matplotlib
pip install -e ".[extras]"    # scikit-learn, PyYAML
pip install -e ".[all]"       # all optional dependencies
pip install -e ".[dev]"       # test dependencies

import torch
from torch_geometric.data import Data
from kgcnn_torch.models.schnet import SchNetModel

model = SchNetModel(
    node_dim=64,
    depth=4,
    units=128,
    gauss_bins=20,
    gauss_distance=4.0,
    num_targets=1,
)

data = Data(
    z=torch.tensor([6, 1, 1, 1, 1]),
    pos=torch.randn(5, 3),
    edge_index=torch.tensor(
        [[0, 0, 0, 0, 1, 2, 3, 4], [1, 2, 3, 4, 0, 0, 0, 0]]
    ),
    batch=torch.zeros(5, dtype=torch.long),
)

prediction = model(data)

from kgcnn_torch.models.force import EnergyForceModel
from kgcnn_torch.models.schnet import SchNetModel

energy_model = SchNetModel(node_dim=128, depth=6, units=128, num_targets=1)
model = EnergyForceModel(
    energy_model=energy_model,
    coordinate_input="pos",
    output_as_dict=True,
    is_physical_force=True,  # force = -dE/dR
)

data.pos.requires_grad_(True)
result = model(data)
energy, forces = result["energy"], result["force"]

Training scripts use JSON hyperparameter configurations in training_scripts/hyper.

# Graph-level property prediction
python training_scripts/train_graph.py \
  --hyper training_scripts/hyper/hyper_esol.json \
  --category SchNet \
  --device cuda \
  --output results/

# Energy-and-force prediction
python training_scripts/train_force.py \
  --hyper training_scripts/hyper/hyper_md17_revised.json \
  --category SchNet \
  --output results/

See notebooks for end-to-end tutorials and CONVERSION_GUIDE.md for implementation and API differences from the Keras project.

kgcnn_torch/
  data/              Dataset loaders, graph containers, and transforms
  layers/            Message passing, geometry, pooling, and neural layers
  models/            GNN architectures and the energy/force wrapper
  training/          Trainer, callbacks, schedulers, metrics, and configuration
  molecule/          Molecule conversion and preprocessing helpers
  crystal/           Crystal and periodic-structure helpers
training_scripts/    Command-line training entry points and configurations
notebooks/           Tutorials and workflow examples
tests/               Unit tests

After installing development dependencies, run:

pytest tests

Some top-level training and dataset checks download data or require optional scientific packages; they are intentionally not part of the minimal unit-test command above.

This project was developed by porting and adapting concepts, model configurations, and parts of the data-processing approach from kgcnn, the Keras graph-convolution library maintained by aimat-lab. We thank its authors and contributors for their foundational open-source work. Please cite and acknowledge the original project when this port is used in research; see the upstream project's citation guidance.

The upstream kgcnn project is distributed under the MIT License. This repository retains the required upstream copyright notice in LICENSE.

kgcnn-torch is distributed under the MIT License. It includes adapted work from aimat-lab/kgcnn; see the license file and acknowledgements above for attribution.