From eacfdfb14b5779277e0fe4130fdf94820e689dce Mon Sep 17 00:00:00 2001 From: CoreLeader Date: Wed, 15 Dec 2021 17:30:00 +0800 Subject: [PATCH] Implement decoupled encoder and decoder for PYG Backend --- autogl/module/model/decoders/_dgl/__init__.py | 5 + .../model/decoders/_dgl/_dgl_decoders.py | 6 +- .../model/decoders/_pyg/_pyg_decoders.py | 134 +++++++++++++++ autogl/module/model/decoders/base_decoder.py | 4 +- autogl/module/model/encoders/_pyg/_gat.py | 155 ++++++++++++++++++ autogl/module/model/encoders/_pyg/_gcn.py | 108 ++++++++++++ autogl/module/model/encoders/_pyg/_gin.py | 155 ++++++++++++++++++ autogl/module/model/encoders/_pyg/_sage.py | 114 +++++++++++++ autogl/module/model/encoders/base_encoder.py | 4 +- 9 files changed, 680 insertions(+), 5 deletions(-) create mode 100644 autogl/module/model/decoders/_pyg/_pyg_decoders.py create mode 100644 autogl/module/model/encoders/_pyg/_gat.py create mode 100644 autogl/module/model/encoders/_pyg/_gcn.py create mode 100644 autogl/module/model/encoders/_pyg/_gin.py create mode 100644 autogl/module/model/encoders/_pyg/_sage.py diff --git a/autogl/module/model/decoders/_dgl/__init__.py b/autogl/module/model/decoders/_dgl/__init__.py index e69de29..f71a686 100644 --- a/autogl/module/model/decoders/_dgl/__init__.py +++ b/autogl/module/model/decoders/_dgl/__init__.py @@ -0,0 +1,5 @@ +from ._dgl_decoders import ( + LogSoftmaxDecoderMaintainer, + GINDecoderMaintainer, + TopKDecoderMaintainer +) diff --git a/autogl/module/model/decoders/_dgl/_dgl_decoders.py b/autogl/module/model/decoders/_dgl/_dgl_decoders.py index 737e2de..3824504 100644 --- a/autogl/module/model/decoders/_dgl/_dgl_decoders.py +++ b/autogl/module/model/decoders/_dgl/_dgl_decoders.py @@ -17,8 +17,10 @@ class _LogSoftmaxDecoder(torch.nn.Module): return torch.nn.functional.log_softmax(features[-1], dim=-1) -@decoder_registry.DecoderUniversalRegistry.register_decoder('gin') -@decoder_registry.DecoderUniversalRegistry.register_decoder('gin_decoder') +@decoder_registry.DecoderUniversalRegistry.register_decoder('log_softmax') +@decoder_registry.DecoderUniversalRegistry.register_decoder('log_softmax_decoder') +@decoder_registry.DecoderUniversalRegistry.register_decoder('LogSoftmax'.lower()) +@decoder_registry.DecoderUniversalRegistry.register_decoder('LogSoftmax_decoder'.lower()) class LogSoftmaxDecoderMaintainer(base_decoder.BaseAutoDecoderMaintainer): def _initialize(self, encoder, *args, **kwargs) -> _typing.Optional[bool]: self._decoder = _LogSoftmaxDecoder().to(self.device) diff --git a/autogl/module/model/decoders/_pyg/_pyg_decoders.py b/autogl/module/model/decoders/_pyg/_pyg_decoders.py new file mode 100644 index 0000000..433373e --- /dev/null +++ b/autogl/module/model/decoders/_pyg/_pyg_decoders.py @@ -0,0 +1,134 @@ +import torch.nn.functional +import typing as _typing +import torch_geometric +from torch_geometric.nn.glob import global_add_pool +from ...encoders import base_encoder +from .. import base_decoder, decoder_registry +from ... import _utils + + +class _LogSoftmaxDecoder(torch.nn.Module): + def forward(self, features: _typing.Sequence[torch.Tensor], *__args, **__kwargs) -> torch.Tensor: + return torch.nn.functional.log_softmax(features[-1]) + + +@decoder_registry.DecoderUniversalRegistry.register_decoder('log_softmax') +@decoder_registry.DecoderUniversalRegistry.register_decoder('log_softmax_decoder') +@decoder_registry.DecoderUniversalRegistry.register_decoder('LogSoftmax'.lower()) +@decoder_registry.DecoderUniversalRegistry.register_decoder('LogSoftmax_decoder'.lower()) +class LogSoftmaxDecoderMaintainer(base_decoder.BaseAutoDecoderMaintainer): + def _initialize(self, *args, **kwargs) -> _typing.Optional[bool]: + self._decoder = _LogSoftmaxDecoder().to(self.device) + return True + + +class _GINDecoder(torch.nn.Module): + def __init__( + self, _final_dimension: int, hidden_dimension: int, output_dimension: int, + _act: _typing.Optional[str], _dropout: _typing.Optional[float], + num_graph_features: _typing.Optional[int] + ): + super(_GINDecoder, self).__init__() + if ( + isinstance(num_graph_features, int) + and num_graph_features > 0 + ): + _final_dimension += num_graph_features + self.__num_graph_features: _typing.Optional[int] = num_graph_features + else: + self.__num_graph_features: _typing.Optional[int] = None + self._fc1: torch.nn.Linear = torch.nn.Linear( + _final_dimension, hidden_dimension + ) + self._fc2: torch.nn.Linear = torch.nn.Linear( + hidden_dimension, output_dimension + ) + self._act: _typing.Optional[str] = _act + self._dropout: _typing.Optional[float] = _dropout + + def forward( + self, features: _typing.Sequence[torch.Tensor], + data: torch_geometric.data.Data, *__args, **__kwargs + ): + feature = features[-1] + feature = global_add_pool(feature, data.batch) + if ( + isinstance(self.__num_graph_features, int) + and self.__num_graph_features > 0 + ): + if ( + hasattr(data, 'gf') and + isinstance(data.gf, torch.Tensor) and data.gf.dim() == 2 and + data.gf.size() == (feature.size(0), self.__num_graph_features) + ): + graph_features: torch.Tensor = data.gf + else: + raise ValueError( + f"The provided data is expected to contain property 'gf' " + f"with {self.__num_graph_features} dimensions as graph feature" + ) + feature: torch.Tensor = torch.cat([feature, graph_features], dim=-1) + feature: torch.Tensor = self._fc1(feature) + feature: torch.Tensor = _utils.activation.activation_func(feature, self._act) + if isinstance(self._dropout, float) and 0 <= self._dropout <= 1: + feature: torch.Tensor = torch.nn.functional.dropout( + feature, self._dropout, self.training + ) + feature: torch.Tensor = self._fc2(feature) + return torch.nn.functional.log_softmax(feature, dim=-1) + + +@decoder_registry.DecoderUniversalRegistry.register_decoder('GIN'.lower()) +@decoder_registry.DecoderUniversalRegistry.register_decoder('GINPool'.lower()) +@decoder_registry.DecoderUniversalRegistry.register_decoder('GINPool_decoder'.lower()) +class GINDecoderMaintainer(base_decoder.BaseAutoDecoderMaintainer): + def _initialize(self, encoder: base_encoder.AutoHomogeneousEncoderMaintainer, *args, **kwargs) -> _typing.Optional[bool]: + if ( + isinstance(getattr(self, "num_graph_features"), int) and + getattr(self, "num_graph_features") > 0 + ): + num_graph_features: _typing.Optional[int] = getattr(self, "num_graph_features") + else: + num_graph_features: _typing.Optional[int] = None + self._decoder = _GINDecoder( + tuple(encoder.get_output_dimensions())[-1], + self.hyper_parameters['hidden'], self.output_dimension, + self.hyper_parameters['act'], self.hyper_parameters['dropout'], + num_graph_features + ).to(self.device) + return True + + def __init__( + self, output_dimension: _typing.Optional[int] = ..., + device: _typing.Union[torch.device, str, int, None] = ..., + *args, **kwargs + ): + super(GINDecoderMaintainer, self).__init__( + output_dimension, device, *args, **kwargs + ) + self.hyper_parameter_space = [ + { + "parameterName": "hidden", + "type": "INTEGER", + "maxValue": 64, + "minValue": 8, + "scalingType": "LINEAR", + }, + { + "parameterName": "act", + "type": "CATEGORICAL", + "feasiblePoints": ["leaky_relu", "relu", "elu", "tanh"], + }, + { + "parameterName": "dropout", + "type": "DOUBLE", + "maxValue": 0.9, + "minValue": 0.1, + "scalingType": "LINEAR", + } + ] + self.hyper_parameters = { + "hidden": 32, + "act": "relu", + "dropout": 0.5 + } diff --git a/autogl/module/model/decoders/base_decoder.py b/autogl/module/model/decoders/base_decoder.py index 7279f5f..5c12c23 100644 --- a/autogl/module/model/decoders/base_decoder.py +++ b/autogl/module/model/decoders/base_decoder.py @@ -5,7 +5,9 @@ from ..encoders import base_encoder class BaseAutoDecoderMaintainer(AutoModule): - def _initialize(self, encoder, *args, **kwargs) -> _typing.Optional[bool]: + def _initialize( + self, encoder: base_encoder.AutoHomogeneousEncoderMaintainer, *args, **kwargs + ) -> _typing.Optional[bool]: """ Abstract initialization method to override """ raise NotImplementedError diff --git a/autogl/module/model/encoders/_pyg/_gat.py b/autogl/module/model/encoders/_pyg/_gat.py new file mode 100644 index 0000000..301995c --- /dev/null +++ b/autogl/module/model/encoders/_pyg/_gat.py @@ -0,0 +1,155 @@ +import torch.nn.functional +import typing as _typing +import torch_geometric +from torch_geometric.nn.conv import GATConv +from .. import base_encoder, encoder_registry +from ... import _utils + + +class GATUtils: + @classmethod + def to_total_hidden_dimensions( + cls, per_head_output_dimensions: _typing.Sequence[int], + num_hidden_heads: int, num_output_heads: int + ) -> _typing.Sequence[int]: + return [ + d * (num_hidden_heads if layer < (len(per_head_output_dimensions) - 1) else num_output_heads) + for layer, d in enumerate(per_head_output_dimensions) + ] + + +class _GAT(torch.nn.Module): + def __init__( + self, input_dimension: int, + per_head_output_dimensions: _typing.Sequence[int], + num_hidden_heads: int, num_output_heads: int, + _dropout: float, _act: _typing.Optional[str] + ): + super(_GAT, self).__init__() + self._dropout: float = _dropout + self._act: _typing.Optional[str] = _act + total_output_dimensions: _typing.Sequence[int] = ( + GATUtils.to_total_hidden_dimensions( + per_head_output_dimensions, num_hidden_heads, num_output_heads + ) + ) + num_layers = len(per_head_output_dimensions) + self.__convolution_layers: torch.nn.ModuleList = torch.nn.ModuleList() + for layer in range(len(per_head_output_dimensions)): + self.__convolution_layers.append( + GATConv( + input_dimension if layer == 0 else total_output_dimensions[layer - 1], + per_head_output_dimensions[layer], + num_hidden_heads if layer < num_layers - 1 else num_output_heads, + dropout=_dropout + ) + ) + + def forward(self, data: torch_geometric.data.Data, *__args, **__kwargs): + x: torch.Tensor = data.x + edge_index: torch.LongTensor = data.edge_index + if ( + isinstance(getattr(data, "edge_weight"), torch.Tensor) + and torch.is_tensor(data.edge_weight) + ): + edge_weight: _typing.Optional[torch.Tensor] = data.edge_weight + else: + edge_weight: _typing.Optional[torch.Tensor] = None + results: _typing.MutableSequence[torch.Tensor] = [] + for layer, _gat in enumerate(self.__convolution_layers): + x: torch.Tensor = torch.nn.functional.dropout( + x, self._dropout, self.training + ) + x: torch.Tensor = _gat(x, edge_index, edge_weight) + if layer < len(self.__convolution_layers) - 1: + x: torch.Tensor = _utils.activation.activation_func(x, self._act) + results.append(x) + return results + + +@encoder_registry.EncoderUniversalRegistry.register_encoder('gat') +@encoder_registry.EncoderUniversalRegistry.register_encoder('gat_encoder') +class GATEncoderMaintainer(base_encoder.AutoHomogeneousEncoderMaintainer): + def _initialize(self) -> _typing.Optional[bool]: + dimensions = list(self.hyper_parameters['hidden']) + if ( + self.final_dimension not in (Ellipsis, None) + and isinstance(self.final_dimension, int) + and self.final_dimension > 0 + ): + dimensions.append(self.final_dimension) + self._encoder = _GAT( + self.input_dimension, self.hyper_parameters['hidden'], + self.hyper_parameters['num_hidden_heads'], + self.hyper_parameters['num_output_heads'], + self.hyper_parameters['dropout'], + self.hyper_parameters['act'] + ) + return True + + def get_output_dimensions(self) -> _typing.Iterable[int]: + temp = list(self.hyper_parameters["hidden"]) + if ( + self.final_dimension not in (Ellipsis, None) and + isinstance(self.final_dimension, int) and + self.final_dimension > 0 + ): + temp.append(self.final_dimension) + return GATUtils.to_total_hidden_dimensions( + temp, + self.hyper_parameters['num_hidden_heads'], + self.hyper_parameters['num_output_heads'] + ) + + def __init__( + self, + input_dimension: _typing.Optional[int] = ..., + final_dimension: _typing.Optional[int] = ..., + device: _typing.Union[torch.device, str, int, None] = ..., + *args, **kwargs + ): + super(GATEncoderMaintainer, self).__init__( + input_dimension, final_dimension, device, *args, **kwargs + ) + self.hyper_parameter_space = [ + { + "parameterName": "num_layers", + "type": "DISCRETE", + "feasiblePoints": "2,3,4", + }, + { + "parameterName": "hidden", + "type": "NUMERICAL_LIST", + "numericalType": "INTEGER", + "length": 3, + "minValue": [8, 8, 8], + "maxValue": [64, 64, 64], + "scalingType": "LOG", + "cutPara": ("num_layers",), + "cutFunc": lambda x: x[0] - 1, + }, + { + "parameterName": "dropout", + "type": "DOUBLE", + "maxValue": 0.8, + "minValue": 0.2, + "scalingType": "LINEAR", + }, + { + "parameterName": "heads", + "type": "DISCRETE", + "feasiblePoints": "2,4,8,16", + }, + { + "parameterName": "act", + "type": "CATEGORICAL", + "feasiblePoints": ["leaky_relu", "relu", "elu", "tanh"], + }, + ] + self.hyper_parameters = { + "num_layers": 2, + "hidden": [32], + "heads": 4, + "dropout": 0.2, + "act": "leaky_relu", + } diff --git a/autogl/module/model/encoders/_pyg/_gcn.py b/autogl/module/model/encoders/_pyg/_gcn.py new file mode 100644 index 0000000..da00f75 --- /dev/null +++ b/autogl/module/model/encoders/_pyg/_gcn.py @@ -0,0 +1,108 @@ +import torch.nn.functional +import typing as _typing +import torch_geometric +from torch_geometric.nn.conv import GCNConv +from .. import base_encoder, encoder_registry +from ... import _utils + + +class _GCN(torch.nn.Module): + def __init__( + self, input_dimension: int, dimensions: _typing.Sequence[int], + _act: _typing.Optional[str], _dropout: _typing.Optional[float] + ): + super(_GCN, self).__init__() + self._act: _typing.Optional[str] = _act + self._dropout: _typing.Optional[float] = _dropout + self.__convolution_layers: torch.nn.ModuleList = torch.nn.ModuleList() + for layer, output_dimension in enumerate(dimensions): + self.__convolution_layers.append( + GCNConv(input_dimension if layer == 0 else dimensions[layer - 1], output_dimension) + ) + + def forward( + self, data: torch_geometric.data.Data, *__args, **__kwargs + ) -> _typing.Sequence[torch.Tensor]: + x: torch.Tensor = data.x + edge_index: torch.LongTensor = data.edge_index + if ( + isinstance(getattr(data, "edge_weight"), torch.Tensor) + and torch.is_tensor(data.edge_weight) + ): + edge_weight: _typing.Optional[torch.Tensor] = data.edge_weight + else: + edge_weight: _typing.Optional[torch.Tensor] = None + results: _typing.MutableSequence[torch.Tensor] = [] + for layer, convolution_layer in enumerate(self.__convolution_layers): + x = convolution_layer(x, edge_index, edge_weight) + if layer < len(self.__convolution_layers) - 1: + x: torch.Tensor = _utils.activation.activation_func(x, self._act) + if isinstance(self._dropout, float) and 0 <= self._dropout <= 1: + x = torch.nn.functional.dropout(x, self._dropout, self.training) + results.append(x) + return results + + +@encoder_registry.EncoderUniversalRegistry.register_encoder('gcn') +@encoder_registry.EncoderUniversalRegistry.register_encoder('gcn_encoder') +class GCNEncoderMaintainer(base_encoder.AutoHomogeneousEncoderMaintainer): + def _initialize(self) -> _typing.Optional[bool]: + dimensions = list(self.hyper_parameters['hidden']) + if ( + self.final_dimension not in (Ellipsis, None) + and isinstance(self.final_dimension, int) + and self.final_dimension > 0 + ): + dimensions.append(self.final_dimension) + self._encoder = _GCN( + self.input_dimension, dimensions, + self.hyper_parameters['act'], self.hyper_parameters['dropout'] + ) + return True + + def __init__( + self, + input_dimension: _typing.Optional[int] = ..., + final_dimension: _typing.Optional[int] = ..., + device: _typing.Union[torch.device, str, int, None] = ..., + *args, **kwargs + ): + super(GCNEncoderMaintainer, self).__init__( + input_dimension, final_dimension, device, *args, **kwargs + ) + self.hyper_parameter_space = [ + { + "parameterName": "num_layers", + "type": "DISCRETE", + "feasiblePoints": "2,3,4", + }, + { + "parameterName": "hidden", + "type": "NUMERICAL_LIST", + "numericalType": "INTEGER", + "length": 3, + "minValue": [8, 8, 8], + "maxValue": [128, 128, 128], + "scalingType": "LOG", + "cutPara": ("num_layers",), + "cutFunc": lambda x: x[0] - 1, + }, + { + "parameterName": "dropout", + "type": "DOUBLE", + "maxValue": 0.8, + "minValue": 0.2, + "scalingType": "LINEAR", + }, + { + "parameterName": "act", + "type": "CATEGORICAL", + "feasiblePoints": ["leaky_relu", "relu", "elu", "tanh"], + }, + ] + self.hyper_parameters = { + "num_layers": 2, + "hidden": [16], + "dropout": 0.2, + "act": "leaky_relu", + } diff --git a/autogl/module/model/encoders/_pyg/_gin.py b/autogl/module/model/encoders/_pyg/_gin.py new file mode 100644 index 0000000..fccc8fe --- /dev/null +++ b/autogl/module/model/encoders/_pyg/_gin.py @@ -0,0 +1,155 @@ +import typing as _typing +import torch.nn.functional +import torch_geometric +from torch_geometric.nn.conv import GINConv +from .. import base_encoder, encoder_registry +from ... import _utils + + +class _GIN(torch.nn.Module): + def __init__( + self, input_dimension: int, + dimensions: _typing.Sequence[int], + _act: str, _dropout: float, + mlp_layers: int, _eps: str + ): + super(_GIN, self).__init__() + + self._act: str = _act + + def _get_act() -> torch.nn.Module: + if _act == 'leaky_relu': + return torch.nn.LeakyReLU() + elif _act == 'relu': + return torch.nn.ReLU() + elif _act == 'elu': + return torch.nn.ELU() + elif _act == 'tanh': + return torch.nn.Tanh() + elif _act == 'PReLU'.lower(): + return torch.nn.PReLU() + else: + return torch.nn.ReLU() + + convolution_layers: torch.nn.ModuleList = torch.nn.ModuleList() + batch_normalizations: torch.nn.ModuleList = torch.nn.ModuleList() + + __mlp_layers = [torch.nn.Linear(input_dimension, dimensions[0])] + for _ in range(mlp_layers - 1): + __mlp_layers.append(_get_act()) + __mlp_layers.append(torch.nn.Linear(dimensions[0], dimensions[0])) + convolution_layers.append( + GINConv(torch.nn.Sequential(*__mlp_layers), train_eps=_eps == "True") + ) + batch_normalizations.append(torch.nn.BatchNorm1d(dimensions[0])) + + num_layers: int = len(dimensions) + for layer in range(num_layers - 1): + __mlp_layers = [torch.nn.Linear(dimensions[layer], dimensions[layer + 1])] + for _ in range(mlp_layers - 1): + __mlp_layers.append(_get_act()) + __mlp_layers.append( + torch.nn.Linear(dimensions[layer + 1], dimensions[layer + 1]) + ) + convolution_layers.append( + GINConv(torch.nn.Sequential(*__mlp_layers), train_eps=_eps == "True") + ) + batch_normalizations.append( + torch.nn.BatchNorm1d(dimensions[layer + 1]) + ) + + self.__convolution_layers: torch.nn.ModuleList = convolution_layers + self.__batch_normalizations: torch.nn.ModuleList = batch_normalizations + + def forward( + self, data: torch_geometric.data.Data, *__args, **__kwargs + ) -> _typing.Sequence[torch.Tensor]: + x: torch.Tensor = data.x + edge_index: torch.Tensor = data.edge_index + + results: _typing.MutableSequence[torch.Tensor] = [] + num_layers = len(self.__convolution_layers) + for layer in range(num_layers): + x: torch.Tensor = self.__convolution_layers[layer](x, edge_index) + x: torch.Tensor = _utils.activation.activation_func(x, self._act) + x: torch.Tensor = self.__batch_normalizations[layer](x) + results.append(x) + return results + + +class GINEncoderMaintainer(base_encoder.AutoHomogeneousEncoderMaintainer): + def _initialize(self) -> _typing.Optional[bool]: + dimensions = list(self.hyper_parameters['hidden']) + if ( + self.final_dimension not in (Ellipsis, None) + and isinstance(self.final_dimension, int) + and self.final_dimension > 0 + ): + dimensions.append(self.final_dimension) + self._encoder = _GIN( + self.input_dimension, dimensions, + self.hyper_parameters['act'], + self.hyper_parameters['dropout'], + self.hyper_parameters['mlp_layers'], + self.hyper_parameters['eps'] + ).to(self.device) + return True + + def __init__( + self, + input_dimension: _typing.Optional[int] = ..., + final_dimension: _typing.Optional[int] = ..., + device: _typing.Union[torch.device, str, int, None] = ..., + *args, **kwargs + ): + super(GINEncoderMaintainer, self).__init__( + input_dimension, final_dimension, device, *args, **kwargs + ) + self.hyper_parameter_space = [ + { + "parameterName": "num_layers", + "type": "DISCRETE", + "feasiblePoints": "4,5,6", + }, + { + "parameterName": "hidden", + "type": "NUMERICAL_LIST", + "numericalType": "INTEGER", + "length": 5, + "minValue": [8, 8, 8, 8, 8], + "maxValue": [64, 64, 64, 64, 64], + "scalingType": "LOG", + "cutPara": ("num_layers",), + "cutFunc": lambda x: x[0] - 1, + }, + { + "parameterName": "dropout", + "type": "DOUBLE", + "maxValue": 0.9, + "minValue": 0.1, + "scalingType": "LINEAR", + }, + { + "parameterName": "act", + "type": "CATEGORICAL", + "feasiblePoints": ["leaky_relu", "relu", "elu", "tanh"], + }, + { + "parameterName": "eps", + "type": "CATEGORICAL", + "feasiblePoints": ["True", "False"], + }, + { + "parameterName": "mlp_layers", + "type": "DISCRETE", + "feasiblePoints": "2,3,4", + }, + ] + self.hyper_parameters = { + "num_layers": 3, + "hidden": [64, 32], + "dropout": 0.5, + "act": "relu", + "eps": "True", + "mlp_layers": 2, + } diff --git a/autogl/module/model/encoders/_pyg/_sage.py b/autogl/module/model/encoders/_pyg/_sage.py new file mode 100644 index 0000000..8d54d4e --- /dev/null +++ b/autogl/module/model/encoders/_pyg/_sage.py @@ -0,0 +1,114 @@ +import torch.nn.functional +import typing as _typing +import torch_geometric +from torch_geometric.nn.conv import SAGEConv +from .. import base_encoder, encoder_registry +from ... import _utils + + +class _SAGE(torch.nn.Module): + def __init__( + self, input_dimension: int, dimensions: _typing.Sequence[int], + _act: _typing.Optional[str], _dropout: _typing.Optional[float], + aggr: _typing.Optional[str] + ): + super(_SAGE, self).__init__() + self._act: _typing.Optional[str] = _act + self._dropout: _typing.Optional[float] = _dropout + self.__convolution_layers: torch.nn.ModuleList = torch.nn.ModuleList() + for layer, output_dimension in enumerate(dimensions): + self.__convolution_layers.append( + SAGEConv( + input_dimension if layer == 0 else dimensions[layer - 1], output_dimension, + aggr=aggr + ) + ) + + def forward( + self, data: torch_geometric.data.Data, *__args, **__kwargs + ) -> _typing.Sequence[torch.Tensor]: + x: torch.Tensor = data.x + edge_index: torch.LongTensor = data.edge_index + results: _typing.MutableSequence[torch.Tensor] = [] + for layer, convolution_layer in enumerate(self.__convolution_layers): + x = convolution_layer(x, edge_index) + if layer < len(self.__convolution_layers) - 1: + x = _utils.activation.activation_func(x, self._act) + if isinstance(self._dropout, float) and 0 <= self._dropout <= 1: + x = torch.nn.functional.dropout(x, self._dropout, self.training) + results.append(x) + return results + + +@encoder_registry.EncoderUniversalRegistry.register_encoder('sage') +@encoder_registry.EncoderUniversalRegistry.register_encoder('sage_encoder') +@encoder_registry.EncoderUniversalRegistry.register_encoder('GraphSAGE'.lower()) +@encoder_registry.EncoderUniversalRegistry.register_encoder('GraphSAGE_encoder'.lower()) +class SAGEEncoderMaintainer(base_encoder.AutoHomogeneousEncoderMaintainer): + def _initialize(self) -> _typing.Optional[bool]: + dimensions = list(self.hyper_parameters['hidden']) + if ( + self.final_dimension not in (Ellipsis, None) + and isinstance(self.final_dimension, int) + and self.final_dimension > 0 + ): + dimensions.append(self.final_dimension) + self._encoder = _SAGE( + self.input_dimension, dimensions, + self.hyper_parameters['act'], self.hyper_parameters['dropout'], + self.hyper_parameters['agg'] + ) + return True + + def __init__( + self, + input_dimension: _typing.Optional[int] = ..., + final_dimension: _typing.Optional[int] = ..., + device: _typing.Union[torch.device, str, int, None] = ..., + *args, **kwargs + ): + super(SAGEEncoderMaintainer, self).__init__( + input_dimension, final_dimension, device, *args, **kwargs + ) + self.hyper_parameter_space = [ + { + "parameterName": "num_layers", + "type": "DISCRETE", + "feasiblePoints": "2,3,4", + }, + { + "parameterName": "hidden", + "type": "NUMERICAL_LIST", + "numericalType": "INTEGER", + "length": 3, + "minValue": [8, 8, 8], + "maxValue": [128, 128, 128], + "scalingType": "LOG", + "cutPara": ("num_layers",), + "cutFunc": lambda x: x[0] - 1, + }, + { + "parameterName": "dropout", + "type": "DOUBLE", + "maxValue": 0.8, + "minValue": 0.2, + "scalingType": "LINEAR", + }, + { + "parameterName": "act", + "type": "CATEGORICAL", + "feasiblePoints": ["leaky_relu", "relu", "elu", "tanh"], + }, + { + "parameterName": "agg", + "type": "CATEGORICAL", + "feasiblePoints": ["mean", "add", "max"], + }, + ] + self.hyper_parameters = { + "num_layers": 3, + "hidden": [64, 32], + "dropout": 0.5, + "act": "relu", + "agg": "mean", + } diff --git a/autogl/module/model/encoders/base_encoder.py b/autogl/module/model/encoders/base_encoder.py index 4c47d06..0e9d23c 100644 --- a/autogl/module/model/encoders/base_encoder.py +++ b/autogl/module/model/encoders/base_encoder.py @@ -77,8 +77,8 @@ class AutoHomogeneousEncoderMaintainer(BaseAutoEncoderMaintainer): new_kwargs = dict(self.__kwargs) new_kwargs.update(kwargs) duplicate: AutoHomogeneousEncoderMaintainer = self.__class__( - self.input_dimension, self.final_dimension, - False, self.device, **new_kwargs + self.input_dimension, self.final_dimension, self.device, + **new_kwargs ) hp = dict(self.hyper_parameters) hp.update(hyper_parameter)