gradnet 0.4.0

Trainable graph adjacency parameterizations with ODE integration and Lightning training helpers.

GradNet

GradNet is a PyTorch-based framework for AI-enabled optimization of networks. Define static or dynamical objectives and constraints, then discover the optimal network structures.

It encodes the network structure as a differentiable object with optional budget and structure constraints. It lets the users directly optimize static objectives using a lightweight PyTorch Lightning training loop. Alternatively, built-in ODE solvers can be used to define and optimize dynamical objectives.

Illustration of the gradient-based optimization pipeline for network structure.

A random network rewires itself with GradNet to optimize synchronization in the Kuramoto model.

Highlights

• Learn dense or sparse adjacency updates with norm, sign, and symmetry constraints.
• Projected parameterizations that stay differentiable and GPU friendly.
• Torchdiffeq-backed integration utilities for graph-driven dynamical systems.
• Minimal Lightning trainer that wraps loss functions in just a few lines.

Installation

Install the released package from PyPI:

pip install gradnet

To work off the latest sources instead, clone the repository and install in editable mode:

pip install -e .

GradNet targets Python 3.10+ and requires pip 21.3+ (run pip install --upgrade pip if needed). It depends on PyTorch, PyTorch Lightning, torchdiffeq, NumPy, and tqdm (installed automatically by the command above). Install the optional NetworkX helpers with pip install 'gradnet[networkx]' when you need conversions to networkx graphs or plotting utilities that rely on it.

Documentation

Full API documentation, tutorials, and background material live at gradnet.readthedocs.io.

Quickstart

The examples folder contains several examples of how to use GradNet.

Spectral optimization (algebraic connectivity)

Demonstrates a simple example of Configuring a GradNet object restricted to a grid lattice. It defines a simple static loss function (the algebraic connectivity). Then it uses fit to optimize the network structure, all in the first code cell of the notebook. The rest of the notebook is analysis of the optimal grid and comparison of dense and sparse backends.

Kuramoto network optimization

A simple example of dynamical loss and usage of integrate_ode. Demonstrates structural optimization emergent sparsity with no mask.

Zachary's karate club

An xample showing how to optimally modify existing networks.

Modules at a glance

• gradnet.GradNet: wraps dense and sparse parameterizations, supports directed/undirected networks, masking, custom edge-building costs etc.
• gradnet.integrate_ode: torchdiffeq-powered solver with adjoint and event support for adjacency-dependent dynamics.
• gradnet.fit: PyTorch Lightning loop that optimizes a GradNet using user-supplied loss functions.
• gradnet.utils: various helpers functions.

Credits

GradNet relies on (and is inspired by) the following open-source projects:

• PyTorch
• PyTorch Lightning
• torchdiffeq

License

GradNet is released under the BSD 3-Clause License. See LICENSE for details.