From 03335fcd17b516994edbd9198106524ef84f17f8 Mon Sep 17 00:00:00 2001 From: Jorge Prado Date: Mon, 19 Aug 2024 18:15:25 +0200 Subject: [PATCH 1/5] Particlenet added and all km3net dependencies are deleted --- src/graphnet/models/gnn/ParticleNeT.py | 244 +++++++++++++++++++++++++ src/graphnet/models/gnn/__init__.py | 1 + 2 files changed, 245 insertions(+) create mode 100644 src/graphnet/models/gnn/ParticleNeT.py diff --git a/src/graphnet/models/gnn/ParticleNeT.py b/src/graphnet/models/gnn/ParticleNeT.py new file mode 100644 index 000000000..b68f82322 --- /dev/null +++ b/src/graphnet/models/gnn/ParticleNeT.py @@ -0,0 +1,244 @@ +"""Implementation of the ParticleNet GNN model architecture.""" +from typing import List, Optional, Callable, Tuple, Union + +import torch +from torch import Tensor, LongTensor +from torch_geometric.data import Data +from torch_scatter import scatter_max, scatter_mean, scatter_min, scatter_sum + +from graphnet.models.components.layers import DynEdgeConv +from graphnet.models.gnn.gnn import GNN + +GLOBAL_POOLINGS = { + "min": scatter_min, + "max": scatter_max, + "sum": scatter_sum, + "mean": scatter_mean, +} + + +class ParticleNeT(GNN): + """ParticleNeT (dynamical edge convolutional) model. Inspired by: https://arxiv.org/abs/1902.08570""" + + def __init__( + self, + nb_inputs: int, + *, + nb_neighbours: int = 16, + features_subset: Optional[Union[List[int], slice]] = None, + dynamic: bool = True, + dynedge_layer_sizes: Optional[List[Tuple[int, ...]]] = [(64,64,64), (128,128,128), (256,256,256)], + readout_layer_sizes: Optional[List[int]] = [256], + global_pooling_schemes: Optional[Union[str, List[str]]] = "mean", + activation_layer: Optional[str] = "relu", + add_batchnorm_layer: bool = True, + dropout_readout: float = 0.1, + skip_readout: bool = False, + ): + """Construct `ParticleNeT`. + + Args: + nb_inputs: Number of input features on each node. + nb_neighbours: Number of neighbours to used in the k-nearest + neighbour clustering which is performed after each (dynamical) + edge convolution. + features_subset: The subset of latent features on each node that + are used as metric dimensions when performing the k-nearest + neighbours clustering. Defaults to [0,1,2]. + dynamic: wether or not update the edges after every `DynEdgeConv` + block. + dynedge_layer_sizes: The layer sizes, or latent feature dimenions, + used in the `DynEdgeConv` layer. Each entry in + `dynedge_layer_sizes` corresponds to a single `DynEdgeConv` + layer; the integers in the corresponding tuple corresponds to + the layer sizes in the multi-layer perceptron (MLP) that is + applied within each `DynEdgeConv` layer. That is, a list of + size-three tuples means that all `DynEdgeConv` layers contain + a three-layer MLP. + Defaults to [(64, 64, 64), (128, 128, 128), (256, 256, 256)]. + readout_layer_sizes: Hidden layer size in the MLP following the + post-processing _and_ optional global pooling. As this is the + last layer in the model, it yields the output of the `DynEdge` + model. Defaults to [256,]. + global_pooling_schemes: The list global pooling schemes to use. + Options are: "min", "max", "mean", and "sum". + Default to "mean". + activation_layer: The activation function to use in the model. + Default to "relu". + add_batchnorm_layer: Whether to add a batch normalization layer after + each linear layer. Default to True. + dropout_readout: Dropout value to use in the readout layer(s). + Default to 0.1. + skip_readout: Whether to skip the readout layer(s). If `True`, the + output of the last DynEdgeConv block is returned directly. + """ + # Latent feature subset for computing nearest neighbours in ParticleNeT. + if features_subset is None: + features_subset = slice(0, 3) + + # DynEdge layer sizes + if dynedge_layer_sizes is None: + dynedge_layer_sizes = [ + ( + 64, 64, 64 + ), + ( + 128, 128, 128, + ), + ( + 256, 256, 256, + ), + ] + + dynedge_layer_sizes_check = [] + for sizes in dynedge_layer_sizes: + if isinstance(sizes, list): + sizes = tuple(sizes) + dynedge_layer_sizes_check.append(sizes) + + assert isinstance(dynedge_layer_sizes_check, list) + assert len(dynedge_layer_sizes_check) + assert all(isinstance(sizes, tuple) for sizes in dynedge_layer_sizes_check) + assert all(len(sizes) > 0 for sizes in dynedge_layer_sizes_check) + assert all( + all(size > 0 for size in sizes) for sizes in dynedge_layer_sizes_check + ) + + self._dynedge_layer_sizes = dynedge_layer_sizes_check + + # Read-out layer sizes + if readout_layer_sizes is None: + readout_layer_sizes = [ + 256, + ] + + assert isinstance(readout_layer_sizes, list) + assert len(readout_layer_sizes) + assert all(size > 0 for size in readout_layer_sizes) + + self._readout_layer_sizes = readout_layer_sizes + + # Global pooling scheme(s) + if isinstance(global_pooling_schemes, str): + global_pooling_schemes = [global_pooling_schemes] + + if isinstance(global_pooling_schemes, list): + for pooling_scheme in global_pooling_schemes: + assert ( + pooling_scheme in GLOBAL_POOLINGS + ), f"Global pooling scheme {pooling_scheme} not supported." + else: + assert global_pooling_schemes is None + + self._global_pooling_schemes = global_pooling_schemes + + if activation_layer is None or activation_layer.lower() == "relu": + activation_layer = torch.nn.ReLU() + elif activation_layer.lower() == "gelu": + activation_layer = torch.nn.GELU() + else: + raise ValueError( + f"Activation layer {activation_layer} not supported." + ) + + # Base class constructor + super().__init__(nb_inputs, self._readout_layer_sizes[-1]) + + # Remaining member variables() + self._activation = activation_layer + self._nb_inputs = nb_inputs + self._nb_neighbours = nb_neighbours + self._features_subset = features_subset + self._dynamic = dynamic + self._add_batchnorm_layer = add_batchnorm_layer + self._dropout_readout = dropout_readout + self._skip_readout = skip_readout + + self._construct_layers() + + # Builds the network + def _construct_layers(self) -> None: + """Construct layers (torch.nn.Modules).""" + + # Convolutional operations + nb_input_features = self._nb_inputs + + self._conv_layers = torch.nn.ModuleList() + nb_latent_features = nb_input_features + for sizes in self._dynedge_layer_sizes: + layers = [] + layer_sizes = [nb_latent_features] + list(sizes) + for ix, (nb_in, nb_out) in enumerate( + zip(layer_sizes[:-1], layer_sizes[1:]) + ): + if ix == 0: + nb_in *= 2 + layers.append(torch.nn.Linear(nb_in, nb_out)) + if self._add_batchnorm_layer: + layers.append(torch.nn.BatchNorm1d(nb_out)) + layers.append(self._activation) + + conv_layer = DynEdgeConv( + torch.nn.Sequential(*layers), + aggr = "mean", + nb_neighbors=self._nb_neighbours, + features_subset=self._features_subset, + ) + self._conv_layers.append(conv_layer) + + nb_latent_features = nb_out + + # Read-out operations + nb_poolings = ( + len(self._global_pooling_schemes) + if self._global_pooling_schemes + else 1 + ) + nb_latent_features = nb_out * nb_poolings + + readout_layers = [] + layer_sizes = [nb_latent_features] + list(self._readout_layer_sizes) + for nb_in, nb_out in zip(layer_sizes[:-1], layer_sizes[1:]): + readout_layers.append(torch.nn.Linear(nb_in, nb_out)) + readout_layers.append(self._activation) + readout_layers.append(torch.nn.Dropout(self._dropout_readout)) + + self._readout = torch.nn.Sequential(*readout_layers) + + def _global_pooling(self, x: Tensor, batch: LongTensor) -> Tensor: + """Perform global pooling.""" + assert self._global_pooling_schemes + pooled = [] + for pooling_scheme in self._global_pooling_schemes: + pooling_fn = GLOBAL_POOLINGS[pooling_scheme] + pooled_x = pooling_fn(x, index=batch, dim=0) + if isinstance(pooled_x, tuple) and len(pooled_x) == 2: + # `scatter_{min,max}`, which return also an argument, vs. + # `scatter_{mean,sum}` + pooled_x, _ = pooled_x + pooled.append(pooled_x) + + return torch.cat(pooled, dim=1) + + def forward(self, data: Data) -> Tensor: + """Apply learnable forward pass.""" + # Convenience variables + x, edge_index, batch = data.x, data.edge_index, data.batch + + # DynEdge-convolutions + for conv_layer in self._conv_layers: + if self._dynamic: + x, edge_index = conv_layer(x, edge_index, batch) + else: + x, _ = conv_layer(x, edge_index, batch) + + # Read-out + if not self._skip_readout: + # (Optional) Global pooling + if self._global_pooling_schemes: + x = self._global_pooling(x, batch=batch) + + # Read-out + x = self._readout(x) + + return x \ No newline at end of file diff --git a/src/graphnet/models/gnn/__init__.py b/src/graphnet/models/gnn/__init__.py index 9b4de8238..27ac80b6b 100644 --- a/src/graphnet/models/gnn/__init__.py +++ b/src/graphnet/models/gnn/__init__.py @@ -6,3 +6,4 @@ from .dynedge_kaggle_tito import DynEdgeTITO from .RNN_tito import RNN_TITO from .icemix import DeepIce +from .ParticleNeT import ParticleNeT \ No newline at end of file From d1f9c3a9f2ad74a107ea2ae84b487310a69aae99 Mon Sep 17 00:00:00 2001 From: Jorge Prado Date: Tue, 20 Aug 2024 10:56:18 +0200 Subject: [PATCH 2/5] Fix style issues --- src/graphnet/models/gnn/ParticleNeT.py | 51 ++++++++++++++++---------- src/graphnet/models/gnn/__init__.py | 2 +- 2 files changed, 32 insertions(+), 21 deletions(-) diff --git a/src/graphnet/models/gnn/ParticleNeT.py b/src/graphnet/models/gnn/ParticleNeT.py index b68f82322..0775cda99 100644 --- a/src/graphnet/models/gnn/ParticleNeT.py +++ b/src/graphnet/models/gnn/ParticleNeT.py @@ -18,7 +18,10 @@ class ParticleNeT(GNN): - """ParticleNeT (dynamical edge convolutional) model. Inspired by: https://arxiv.org/abs/1902.08570""" + """ParticleNeT (dynamical edge convolutional) model. + + Inspired by: https://arxiv.org/abs/1902.08570 + """ def __init__( self, @@ -27,7 +30,11 @@ def __init__( nb_neighbours: int = 16, features_subset: Optional[Union[List[int], slice]] = None, dynamic: bool = True, - dynedge_layer_sizes: Optional[List[Tuple[int, ...]]] = [(64,64,64), (128,128,128), (256,256,256)], + dynedge_layer_sizes: Optional[List[Tuple[int, ...]]] = [ + (64, 64, 64), + (128, 128, 128), + (256, 256, 256), + ], readout_layer_sizes: Optional[List[int]] = [256], global_pooling_schemes: Optional[Union[str, List[str]]] = "mean", activation_layer: Optional[str] = "relu", @@ -53,7 +60,7 @@ def __init__( layer; the integers in the corresponding tuple corresponds to the layer sizes in the multi-layer perceptron (MLP) that is applied within each `DynEdgeConv` layer. That is, a list of - size-three tuples means that all `DynEdgeConv` layers contain + size-three tuples means that all `DynEdgeConv` layers contain a three-layer MLP. Defaults to [(64, 64, 64), (128, 128, 128), (256, 256, 256)]. readout_layer_sizes: Hidden layer size in the MLP following the @@ -65,7 +72,7 @@ def __init__( Default to "mean". activation_layer: The activation function to use in the model. Default to "relu". - add_batchnorm_layer: Whether to add a batch normalization layer after + add_batchnorm_layer: Whether to add a batch normalization layer after each linear layer. Default to True. dropout_readout: Dropout value to use in the readout layer(s). Default to 0.1. @@ -79,29 +86,34 @@ def __init__( # DynEdge layer sizes if dynedge_layer_sizes is None: dynedge_layer_sizes = [ + (64, 64, 64), ( - 64, 64, 64 - ), - ( - 128, 128, 128, + 128, + 128, + 128, ), ( - 256, 256, 256, + 256, + 256, + 256, ), ] - + dynedge_layer_sizes_check = [] for sizes in dynedge_layer_sizes: if isinstance(sizes, list): - sizes = tuple(sizes) + sizes = tuple(sizes) dynedge_layer_sizes_check.append(sizes) assert isinstance(dynedge_layer_sizes_check, list) assert len(dynedge_layer_sizes_check) - assert all(isinstance(sizes, tuple) for sizes in dynedge_layer_sizes_check) + assert all( + isinstance(sizes, tuple) for sizes in dynedge_layer_sizes_check + ) assert all(len(sizes) > 0 for sizes in dynedge_layer_sizes_check) assert all( - all(size > 0 for size in sizes) for sizes in dynedge_layer_sizes_check + all(size > 0 for size in sizes) + for sizes in dynedge_layer_sizes_check ) self._dynedge_layer_sizes = dynedge_layer_sizes_check @@ -140,7 +152,7 @@ def __init__( raise ValueError( f"Activation layer {activation_layer} not supported." ) - + # Base class constructor super().__init__(nb_inputs, self._readout_layer_sizes[-1]) @@ -159,10 +171,9 @@ def __init__( # Builds the network def _construct_layers(self) -> None: """Construct layers (torch.nn.Modules).""" - # Convolutional operations nb_input_features = self._nb_inputs - + self._conv_layers = torch.nn.ModuleList() nb_latent_features = nb_input_features for sizes in self._dynedge_layer_sizes: @@ -177,10 +188,10 @@ def _construct_layers(self) -> None: if self._add_batchnorm_layer: layers.append(torch.nn.BatchNorm1d(nb_out)) layers.append(self._activation) - + conv_layer = DynEdgeConv( torch.nn.Sequential(*layers), - aggr = "mean", + aggr="mean", nb_neighbors=self._nb_neighbours, features_subset=self._features_subset, ) @@ -231,7 +242,7 @@ def forward(self, data: Data) -> Tensor: x, edge_index = conv_layer(x, edge_index, batch) else: x, _ = conv_layer(x, edge_index, batch) - + # Read-out if not self._skip_readout: # (Optional) Global pooling @@ -241,4 +252,4 @@ def forward(self, data: Data) -> Tensor: # Read-out x = self._readout(x) - return x \ No newline at end of file + return x diff --git a/src/graphnet/models/gnn/__init__.py b/src/graphnet/models/gnn/__init__.py index 27ac80b6b..c5e7398b2 100644 --- a/src/graphnet/models/gnn/__init__.py +++ b/src/graphnet/models/gnn/__init__.py @@ -6,4 +6,4 @@ from .dynedge_kaggle_tito import DynEdgeTITO from .RNN_tito import RNN_TITO from .icemix import DeepIce -from .ParticleNeT import ParticleNeT \ No newline at end of file +from .ParticleNeT import ParticleNeT From 97c997a0cdcfa3143950c2b4b831f7e6892a30d5 Mon Sep 17 00:00:00 2001 From: IvanMM27 <112850777+IvanMM27@users.noreply.github.com> Date: Mon, 26 Aug 2024 11:39:41 +0200 Subject: [PATCH 3/5] line lengths corrected and particlenet.py renaming --- src/graphnet/models/gnn/{ParticleNeT.py => particlenet.py} | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) rename src/graphnet/models/gnn/{ParticleNeT.py => particlenet.py} (99%) diff --git a/src/graphnet/models/gnn/ParticleNeT.py b/src/graphnet/models/gnn/particlenet.py similarity index 99% rename from src/graphnet/models/gnn/ParticleNeT.py rename to src/graphnet/models/gnn/particlenet.py index 0775cda99..4ff51691b 100644 --- a/src/graphnet/models/gnn/ParticleNeT.py +++ b/src/graphnet/models/gnn/particlenet.py @@ -72,14 +72,14 @@ def __init__( Default to "mean". activation_layer: The activation function to use in the model. Default to "relu". - add_batchnorm_layer: Whether to add a batch normalization layer after - each linear layer. Default to True. + add_batchnorm_layer: Whether to add a batch normalization layer + after each linear layer. Default to True. dropout_readout: Dropout value to use in the readout layer(s). Default to 0.1. skip_readout: Whether to skip the readout layer(s). If `True`, the output of the last DynEdgeConv block is returned directly. """ - # Latent feature subset for computing nearest neighbours in ParticleNeT. + # Latent feature subset for computing nearest neighbours in model if features_subset is None: features_subset = slice(0, 3) From 7f1c7606fd1fb1d132dc9114c41a919af44d992c Mon Sep 17 00:00:00 2001 From: IvanMM27 <112850777+IvanMM27@users.noreply.github.com> Date: Mon, 26 Aug 2024 11:40:09 +0200 Subject: [PATCH 4/5] renamed particlenet in __init__.py --- src/graphnet/models/gnn/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/graphnet/models/gnn/__init__.py b/src/graphnet/models/gnn/__init__.py index c5e7398b2..8bf94af19 100644 --- a/src/graphnet/models/gnn/__init__.py +++ b/src/graphnet/models/gnn/__init__.py @@ -6,4 +6,4 @@ from .dynedge_kaggle_tito import DynEdgeTITO from .RNN_tito import RNN_TITO from .icemix import DeepIce -from .ParticleNeT import ParticleNeT +from .particlenet import ParticleNeT From 4e8b327c22dfac83d832ad1526907d2fb855468f Mon Sep 17 00:00:00 2001 From: Jorge Prado Date: Mon, 2 Sep 2024 11:55:07 +0200 Subject: [PATCH 5/5] Fix Style issues --- src/graphnet/models/gnn/particlenet.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/graphnet/models/gnn/particlenet.py b/src/graphnet/models/gnn/particlenet.py index 4ff51691b..cf2d00998 100644 --- a/src/graphnet/models/gnn/particlenet.py +++ b/src/graphnet/models/gnn/particlenet.py @@ -72,7 +72,7 @@ def __init__( Default to "mean". activation_layer: The activation function to use in the model. Default to "relu". - add_batchnorm_layer: Whether to add a batch normalization layer + add_batchnorm_layer: Whether to add a batch normalization layer after each linear layer. Default to True. dropout_readout: Dropout value to use in the readout layer(s). Default to 0.1.