From be3c462b2a20558187ec0cfb3f98c8a56c37950e Mon Sep 17 00:00:00 2001 From: Frozenmad Date: Sun, 11 Jul 2021 15:23:40 +0000 Subject: [PATCH] black format --- autogl/datasets/__init__.py | 18 +- autogl/datasets/pyg.py | 2 +- autogl/module/__init__.py | 9 +- autogl/module/feature/__init__.py | 8 +- autogl/module/feature/auto_feature.py | 2 +- autogl/module/feature/base.py | 2 +- autogl/module/feature/generators/__init__.py | 11 +- autogl/module/feature/generators/base.py | 2 +- autogl/module/feature/generators/pyg.py | 2 +- autogl/module/feature/graph/base.py | 2 +- autogl/module/feature/graph/netlsd.py | 2 +- autogl/module/feature/graph/nx.py | 2 +- autogl/module/feature/selectors/base.py | 2 +- autogl/module/hpo/base.py | 4 +- autogl/module/model/base.py | 8 +- autogl/module/model/gcn.py | 190 +++-- autogl/module/model/graph_saint.py | 232 +++--- autogl/module/model/graphsage.py | 106 ++- autogl/module/nas/__init__.py | 6 +- autogl/module/nas/algorithm/__init__.py | 8 +- autogl/module/nas/algorithm/darts.py | 13 +- autogl/module/nas/algorithm/enas.py | 129 ++-- autogl/module/nas/algorithm/random_search.py | 80 +- autogl/module/nas/algorithm/rl.py | 404 +++++++---- autogl/module/nas/estimator/__init__.py | 8 +- autogl/module/nas/estimator/base.py | 6 +- autogl/module/nas/estimator/one_shot.py | 7 +- autogl/module/nas/estimator/train_scratch.py | 33 +- autogl/module/nas/space/__init__.py | 11 +- autogl/module/nas/space/base.py | 70 +- autogl/module/nas/space/graph_nas.py | 106 ++- autogl/module/nas/space/graph_nas_macro.py | 328 ++++++--- autogl/module/nas/space/operation.py | 35 +- autogl/module/nas/space/single_path.py | 20 +- autogl/module/nas/utils.py | 8 +- autogl/module/train/evaluation.py | 46 +- .../module/train/graph_classification_full.py | 19 +- .../module/train/node_classification_full.py | 5 +- .../node_classification_sampled_trainer.py | 684 +++++++++++------- .../sampling/sampler/graphsaint_sampler.py | 52 +- .../layer_dependent_importance_sampler.py | 369 ++++++---- .../sampling/sampler/neighbor_sampler.py | 111 ++- .../sampler/target_dependant_sampler.py | 191 +++-- autogl/solver/base.py | 70 +- autogl/solver/classifier/graph_classifier.py | 6 +- autogl/solver/classifier/node_classifier.py | 42 +- autogl/solver/utils.py | 15 +- autogl/utils/device.py | 13 +- 48 files changed, 2279 insertions(+), 1220 deletions(-) diff --git a/autogl/datasets/__init__.py b/autogl/datasets/__init__.py index 363b5d1..df450d1 100644 --- a/autogl/datasets/__init__.py +++ b/autogl/datasets/__init__.py @@ -178,17 +178,17 @@ __all__ = [ "OGBLcitationDataset", "OGBLwikikgDataset", "OGBLbiokgDataset", - "GatneDataset", - "AmazonDataset", - "TwitterDataset", + "GatneDataset", + "AmazonDataset", + "TwitterDataset", "YouTubeDataset", - "GTNDataset", - "ACM_GTNDataset", - "DBLP_GTNDataset", + "GTNDataset", + "ACM_GTNDataset", + "DBLP_GTNDataset", "IMDB_GTNDataset", - "HANDataset", - "ACM_HANDataset", - "DBLP_HANDataset", + "HANDataset", + "ACM_HANDataset", + "DBLP_HANDataset", "IMDB_HANDataset", "MatlabMatrix", "BlogcatalogDataset", diff --git a/autogl/datasets/pyg.py b/autogl/datasets/pyg.py index df9f480..7a39d3d 100644 --- a/autogl/datasets/pyg.py +++ b/autogl/datasets/pyg.py @@ -10,7 +10,7 @@ from torch_geometric.datasets import ( QM9, Amazon, Coauthor, - Flickr + Flickr, ) from torch_geometric.utils import remove_self_loops from . import register_dataset diff --git a/autogl/module/__init__.py b/autogl/module/__init__.py index 80f94da..02238e7 100644 --- a/autogl/module/__init__.py +++ b/autogl/module/__init__.py @@ -1,11 +1,4 @@ -from . import ( - feature, - model, - train, - hpo, - nas, - ensemble -) +from . import feature, model, train, hpo, nas, ensemble from .ensemble import * from .feature import * diff --git a/autogl/module/feature/__init__.py b/autogl/module/feature/__init__.py index 5af47c5..f2f9c2b 100644 --- a/autogl/module/feature/__init__.py +++ b/autogl/module/feature/__init__.py @@ -35,14 +35,10 @@ from .generators import ( PYGGenerator, PYGLocalDegreeProfile, PYGNormalizeFeatures, - PYGOneHotDegree + PYGOneHotDegree, ) -from .selectors import ( - BaseSelector, - SeFilterConstant, - SeGBDT -) +from .selectors import BaseSelector, SeFilterConstant, SeGBDT from .graph import ( BaseGraph, diff --git a/autogl/module/feature/auto_feature.py b/autogl/module/feature/auto_feature.py index d04c98c..beed033 100644 --- a/autogl/module/feature/auto_feature.py +++ b/autogl/module/feature/auto_feature.py @@ -232,4 +232,4 @@ class AutoFeatureEngineer(BaseFeatureEngineer): gx = gx[:, sel] x = np.concatenate([x, gx], axis=1) data.x = x - return data \ No newline at end of file + return data diff --git a/autogl/module/feature/base.py b/autogl/module/feature/base.py index 9d1e9cc..4b812ba 100644 --- a/autogl/module/feature/base.py +++ b/autogl/module/feature/base.py @@ -169,4 +169,4 @@ class TransformWrapper(BaseFeature): return self def _transform(self, data=None): - return self._func(data) \ No newline at end of file + return self._func(data) diff --git a/autogl/module/feature/generators/__init__.py b/autogl/module/feature/generators/__init__.py index c15a9b3..45fb044 100644 --- a/autogl/module/feature/generators/__init__.py +++ b/autogl/module/feature/generators/__init__.py @@ -2,7 +2,14 @@ from .base import BaseGenerator from .graphlet import GeGraphlet from .eigen import GeEigen from .page_rank import GePageRank -from .pyg import register_pyg, PYGGenerator, pygfunc, PYGLocalDegreeProfile, PYGNormalizeFeatures, PYGOneHotDegree +from .pyg import ( + register_pyg, + PYGGenerator, + pygfunc, + PYGLocalDegreeProfile, + PYGNormalizeFeatures, + PYGOneHotDegree, +) __all__ = [ "BaseGenerator", @@ -14,5 +21,5 @@ __all__ = [ "PYGGenerator", "PYGLocalDegreeProfile", "PYGNormalizeFeatures", - "PYGOneHotDegree" + "PYGOneHotDegree", ] diff --git a/autogl/module/feature/generators/base.py b/autogl/module/feature/generators/base.py index 0ab2024..194d539 100644 --- a/autogl/module/feature/generators/base.py +++ b/autogl/module/feature/generators/base.py @@ -15,4 +15,4 @@ class GeOnehot(BaseGenerator): def _transform(self, data): fe = np.eye(data.x.shape[0]) data.x = np.concatenate([data.x, fe], axis=1) - return data \ No newline at end of file + return data diff --git a/autogl/module/feature/generators/pyg.py b/autogl/module/feature/generators/pyg.py index c3c5e8e..82731b1 100644 --- a/autogl/module/feature/generators/pyg.py +++ b/autogl/module/feature/generators/pyg.py @@ -86,4 +86,4 @@ class PYGOneHotDegree(PYGGenerator): dsc = self.extract(data) data.x = torch.cat([data.x, dsc], dim=1) return data - """ \ No newline at end of file + """ diff --git a/autogl/module/feature/graph/base.py b/autogl/module/feature/graph/base.py index 1daace9..1f3ba21 100644 --- a/autogl/module/feature/graph/base.py +++ b/autogl/module/feature/graph/base.py @@ -16,4 +16,4 @@ class BaseGraph(BaseFeature): data.gf = torch.FloatTensor([[]]) def _postprocess(self, data): - pass \ No newline at end of file + pass diff --git a/autogl/module/feature/graph/netlsd.py b/autogl/module/feature/graph/netlsd.py index 22859e4..205cda1 100644 --- a/autogl/module/feature/graph/netlsd.py +++ b/autogl/module/feature/graph/netlsd.py @@ -25,4 +25,4 @@ class SgNetLSD(BaseGraph): def _transform(self, data): dsc = torch.FloatTensor([netlsd.heat(data.G, *self._args, **self._kwargs)]) data.gf = torch.cat([data.gf, dsc], dim=1) - return data \ No newline at end of file + return data diff --git a/autogl/module/feature/graph/nx.py b/autogl/module/feature/graph/nx.py index 9179b3b..e79f614 100644 --- a/autogl/module/feature/graph/nx.py +++ b/autogl/module/feature/graph/nx.py @@ -180,4 +180,4 @@ class NxIsEulerian(NxGraph): pass -# till algorithms.flows \ No newline at end of file +# till algorithms.flows diff --git a/autogl/module/feature/selectors/base.py b/autogl/module/feature/selectors/base.py index 2553836..7838ffc 100644 --- a/autogl/module/feature/selectors/base.py +++ b/autogl/module/feature/selectors/base.py @@ -12,4 +12,4 @@ class BaseSelector(BaseFeature): def _transform(self, data): if self._sel is not None: data.x = data.x[:, self._sel] - return data \ No newline at end of file + return data diff --git a/autogl/module/hpo/base.py b/autogl/module/hpo/base.py index fdef8a2..180efd6 100644 --- a/autogl/module/hpo/base.py +++ b/autogl/module/hpo/base.py @@ -30,7 +30,9 @@ class BaseHPOptimizer: raise WrongDependedParameterError("The depended parameter does not exist.") for para in config: - if para["type"] in ("NUMERICAL_LIST", "CATEGORICAL_LIST") and para.get("cutPara", None): + if para["type"] in ("NUMERICAL_LIST", "CATEGORICAL_LIST") and para.get( + "cutPara", None + ): self._depend_map[para["parameterName"]] = para if type(para["cutPara"]) == str: get_depended_para(para["cutPara"]) diff --git a/autogl/module/model/base.py b/autogl/module/model/base.py index e251468..99f2c2c 100644 --- a/autogl/module/model/base.py +++ b/autogl/module/model/base.py @@ -302,8 +302,11 @@ class ClassificationModel(_BaseModel): if "initialize" in kwargs: del kwargs["initialize"] super(ClassificationModel, self).__init__( - initialize=init, hyper_parameter_space=hyper_parameter_space, - hyper_parameter=hyper_parameter, device=device, **kwargs + initialize=init, + hyper_parameter_space=hyper_parameter_space, + hyper_parameter=hyper_parameter, + device=device, + **kwargs ) if num_classes != Ellipsis and type(num_classes) == int: self.__num_classes: int = num_classes if num_classes > 0 else 0 @@ -323,6 +326,7 @@ class ClassificationModel(_BaseModel): def __repr__(self) -> str: import yaml + return yaml.dump(self.hyper_parameter) @property diff --git a/autogl/module/model/gcn.py b/autogl/module/model/gcn.py index 6862036..af6671b 100644 --- a/autogl/module/model/gcn.py +++ b/autogl/module/model/gcn.py @@ -14,36 +14,40 @@ LOGGER = get_logger("GCNModel") class GCN(ClassificationSupportedSequentialModel): class _GCNLayer(torch.nn.Module): def __init__( - self, input_channels: int, output_channels: int, - add_self_loops: bool = True, normalize: bool = True, - activation_name: _typing.Optional[str] = ..., - dropout_probability: _typing.Optional[float] = ... + self, + input_channels: int, + output_channels: int, + add_self_loops: bool = True, + normalize: bool = True, + activation_name: _typing.Optional[str] = ..., + dropout_probability: _typing.Optional[float] = ..., ): super().__init__() self._convolution: GCNConv = GCNConv( - input_channels, output_channels, + input_channels, + output_channels, add_self_loops=bool(add_self_loops), - normalize=bool(normalize) + normalize=bool(normalize), ) if ( - activation_name is not Ellipsis and - activation_name is not None and - type(activation_name) == str + activation_name is not Ellipsis + and activation_name is not None + and type(activation_name) == str ): self._activation_name: _typing.Optional[str] = activation_name else: self._activation_name: _typing.Optional[str] = None if ( - dropout_probability is not Ellipsis and - dropout_probability is not None and - type(dropout_probability) == float + dropout_probability is not Ellipsis + and dropout_probability is not None + and type(dropout_probability) == float ): if dropout_probability < 0: dropout_probability = 0 if dropout_probability > 1: dropout_probability = 1 - self._dropout: _typing.Optional[torch.nn.Dropout] = ( - torch.nn.Dropout(dropout_probability) + self._dropout: _typing.Optional[torch.nn.Dropout] = torch.nn.Dropout( + dropout_probability ) else: self._dropout: _typing.Optional[torch.nn.Dropout] = None @@ -51,13 +55,15 @@ class GCN(ClassificationSupportedSequentialModel): def forward(self, data, enable_activation: bool = True) -> torch.Tensor: x: torch.Tensor = getattr(data, "x") edge_index: torch.LongTensor = getattr(data, "edge_index") - edge_weight: _typing.Optional[torch.Tensor] = getattr(data, "edge_weight", None) + edge_weight: _typing.Optional[torch.Tensor] = getattr( + data, "edge_weight", None + ) """ Validate the arguments """ if not type(x) == type(edge_index) == torch.Tensor: raise TypeError if edge_weight is not None and ( - type(edge_weight) != torch.Tensor or - edge_index.size() != (2, edge_weight.size(0)) + type(edge_weight) != torch.Tensor + or edge_index.size() != (2, edge_weight.size(0)) ): edge_weight: _typing.Optional[torch.Tensor] = None @@ -69,15 +75,16 @@ class GCN(ClassificationSupportedSequentialModel): return x def __init__( - self, - num_features: int, - num_classes: int, - hidden_features: _typing.Sequence[int], - activation_name: str, - dropout: _typing.Union[ - _typing.Optional[float], _typing.Sequence[_typing.Optional[float]] - ] = None, - add_self_loops: bool = True, normalize: bool = True + self, + num_features: int, + num_classes: int, + hidden_features: _typing.Sequence[int], + activation_name: str, + dropout: _typing.Union[ + _typing.Optional[float], _typing.Sequence[_typing.Optional[float]] + ] = None, + add_self_loops: bool = True, + normalize: bool = True, ): if isinstance(dropout, _typing.Sequence): if len(dropout) != len(hidden_features) + 1: @@ -97,9 +104,9 @@ class GCN(ClassificationSupportedSequentialModel): dropout = 0 if dropout > 1: dropout = 1 - dropout_list: _typing.Sequence[_typing.Optional[float]] = ( - [dropout for _ in range(len(hidden_features))] + [None] - ) + dropout_list: _typing.Sequence[_typing.Optional[float]] = [ + dropout for _ in range(len(hidden_features)) + ] + [None] elif dropout in (None, Ellipsis, ...): dropout_list: _typing.Sequence[_typing.Optional[float]] = [ None for _ in range(len(hidden_features) + 1) @@ -111,60 +118,86 @@ class GCN(ClassificationSupportedSequentialModel): ) super().__init__() if len(hidden_features) == 0: - self.__sequential_encoding_layers: torch.nn.ModuleList = torch.nn.ModuleList( - ( - self._GCNLayer( - num_features, num_classes, add_self_loops, normalize, - dropout_probability=dropout_list[0] - ), + self.__sequential_encoding_layers: torch.nn.ModuleList = ( + torch.nn.ModuleList( + ( + self._GCNLayer( + num_features, + num_classes, + add_self_loops, + normalize, + dropout_probability=dropout_list[0], + ), + ) ) ) else: - self.__sequential_encoding_layers: torch.nn.ModuleList = torch.nn.ModuleList() - self.__sequential_encoding_layers.append(self._GCNLayer( - num_features, hidden_features[0], add_self_loops, - normalize, activation_name, dropout_list[0] - )) + self.__sequential_encoding_layers: torch.nn.ModuleList = ( + torch.nn.ModuleList() + ) + self.__sequential_encoding_layers.append( + self._GCNLayer( + num_features, + hidden_features[0], + add_self_loops, + normalize, + activation_name, + dropout_list[0], + ) + ) for hidden_feature_index in range(len(hidden_features)): if hidden_feature_index + 1 < len(hidden_features): - self.__sequential_encoding_layers.append(self._GCNLayer( - hidden_features[hidden_feature_index], - hidden_features[hidden_feature_index + 1], - add_self_loops, normalize, activation_name, - dropout_list[hidden_feature_index + 1] - )) + self.__sequential_encoding_layers.append( + self._GCNLayer( + hidden_features[hidden_feature_index], + hidden_features[hidden_feature_index + 1], + add_self_loops, + normalize, + activation_name, + dropout_list[hidden_feature_index + 1], + ) + ) else: - self.__sequential_encoding_layers.append(self._GCNLayer( - hidden_features[hidden_feature_index], num_classes, - add_self_loops, normalize, - dropout_list[-1] - )) + self.__sequential_encoding_layers.append( + self._GCNLayer( + hidden_features[hidden_feature_index], + num_classes, + add_self_loops, + normalize, + dropout_list[-1], + ) + ) @property def sequential_encoding_layers(self) -> torch.nn.ModuleList: return self.__sequential_encoding_layers - def __extract_edge_indexes_and_weights(self, data) -> _typing.Union[ - _typing.Sequence[_typing.Tuple[torch.LongTensor, _typing.Optional[torch.Tensor]]], - _typing.Tuple[torch.LongTensor, _typing.Optional[torch.Tensor]] + def __extract_edge_indexes_and_weights( + self, data + ) -> _typing.Union[ + _typing.Sequence[ + _typing.Tuple[torch.LongTensor, _typing.Optional[torch.Tensor]] + ], + _typing.Tuple[torch.LongTensor, _typing.Optional[torch.Tensor]], ]: def __compose_edge_index_and_weight( - _edge_index: torch.LongTensor, - _edge_weight: _typing.Optional[torch.Tensor] = None + _edge_index: torch.LongTensor, + _edge_weight: _typing.Optional[torch.Tensor] = None, ) -> _typing.Tuple[torch.LongTensor, _typing.Optional[torch.Tensor]]: if type(_edge_index) != torch.Tensor or _edge_index.dtype != torch.int64: raise TypeError if _edge_weight is not None and ( - type(_edge_weight) != torch.Tensor or - _edge_index.size() != (2, _edge_weight.size(0)) + type(_edge_weight) != torch.Tensor + or _edge_index.size() != (2, _edge_weight.size(0)) ): _edge_weight: _typing.Optional[torch.Tensor] = None return _edge_index, _edge_weight if not ( - hasattr(data, "edge_indexes") and - isinstance(getattr(data, "edge_indexes"), _typing.Sequence) and - len(getattr(data, "edge_indexes")) == len(self.__sequential_encoding_layers) + hasattr(data, "edge_indexes") + and isinstance(getattr(data, "edge_indexes"), _typing.Sequence) + and len(getattr(data, "edge_indexes")) + == len(self.__sequential_encoding_layers) ): return __compose_edge_index_and_weight( getattr(data, "edge_index"), getattr(data, "edge_weight", None) @@ -176,14 +209,16 @@ class GCN(ClassificationSupportedSequentialModel): ) if ( - hasattr(data, "edge_weights") and - isinstance(getattr(data, "edge_weights"), _typing.Sequence) and - len(getattr(data, "edge_weights")) == len(self.__sequential_encoding_layers) + hasattr(data, "edge_weights") + and isinstance(getattr(data, "edge_weights"), _typing.Sequence) + and len(getattr(data, "edge_weights")) + == len(self.__sequential_encoding_layers) ): return [ __compose_edge_index_and_weight(_edge_index, _edge_weight) - for _edge_index, _edge_weight - in zip(getattr(data, "edge_indexes"), getattr(data, "edge_weights")) + for _edge_index, _edge_weight in zip( + getattr(data, "edge_indexes"), getattr(data, "edge_weights") + ) ] else: return [ @@ -193,19 +228,22 @@ class GCN(ClassificationSupportedSequentialModel): def cls_encode(self, data) -> torch.Tensor: edge_indexes_and_weights: _typing.Union[ - _typing.Sequence[_typing.Tuple[torch.LongTensor, _typing.Optional[torch.Tensor]]], - _typing.Tuple[torch.LongTensor, _typing.Optional[torch.Tensor]] + _typing.Sequence[ + _typing.Tuple[torch.LongTensor, _typing.Optional[torch.Tensor]] + ], + _typing.Tuple[torch.LongTensor, _typing.Optional[torch.Tensor]], ] = self.__extract_edge_indexes_and_weights(data) - if ( - (not isinstance(edge_indexes_and_weights, tuple)) - and isinstance(edge_indexes_and_weights[0], tuple) + if (not isinstance(edge_indexes_and_weights, tuple)) and isinstance( + edge_indexes_and_weights[0], tuple ): """ edge_indexes_and_weights is sequence of (edge_index, edge_weight) """ - assert len(edge_indexes_and_weights) == len(self.__sequential_encoding_layers) + assert len(edge_indexes_and_weights) == len( + self.__sequential_encoding_layers + ) x: torch.Tensor = getattr(data, "x") for _edge_index_and_weight, gcn in zip( - edge_indexes_and_weights, self.__sequential_encoding_layers + edge_indexes_and_weights, self.__sequential_encoding_layers ): _temp_data = autogl.data.Data(x=x, edge_index=_edge_index_and_weight[0]) _temp_data.edge_weight = _edge_index_and_weight[1] @@ -215,7 +253,9 @@ class GCN(ClassificationSupportedSequentialModel): """ edge_indexes_and_weights is (edge_index, edge_weight) """ x = getattr(data, "x") for gcn in self.__sequential_encoding_layers: - _temp_data = autogl.data.Data(x=x, edge_index=edge_indexes_and_weights[0]) + _temp_data = autogl.data.Data( + x=x, edge_index=edge_indexes_and_weights[0] + ) _temp_data.edge_weight = edge_indexes_and_weights[1] x = gcn(_temp_data) return x @@ -364,5 +404,5 @@ class AutoGCN(BaseModel): self.hyperparams.get("act"), self.hyperparams.get("dropout", None), bool(self.hyperparams.get("add_self_loops", True)), - bool(self.hyperparams.get("normalize", True)) + bool(self.hyperparams.get("normalize", True)), ).to(self.device) diff --git a/autogl/module/model/graph_saint.py b/autogl/module/model/graph_saint.py index 9ebbb04..1b4bc7a 100644 --- a/autogl/module/model/graph_saint.py +++ b/autogl/module/model/graph_saint.py @@ -11,8 +11,12 @@ class _GraphSAINTAggregationLayers: class MultiOrderAggregationLayer(torch.nn.Module): class Order0Aggregator(torch.nn.Module): def __init__( - self, input_dimension: int, output_dimension: int, bias: bool = True, - activation: _typing.Optional[str] = "ReLU", batch_norm: bool = True + self, + input_dimension: int, + output_dimension: int, + bias: bool = True, + activation: _typing.Optional[str] = "ReLU", + batch_norm: bool = True, ): super().__init__() if not type(input_dimension) == type(output_dimension) == int: @@ -21,16 +25,17 @@ class _GraphSAINTAggregationLayers: raise ValueError if not type(bias) == bool: raise TypeError - self.__linear_transform = torch.nn.Linear(input_dimension, output_dimension, bias) + self.__linear_transform = torch.nn.Linear( + input_dimension, output_dimension, bias + ) self.__linear_transform.reset_parameters() if type(activation) == str: if activation.lower() == "ReLU".lower(): self.__activation = torch.nn.functional.relu elif activation.lower() == "elu": self.__activation = torch.nn.functional.elu - elif ( - hasattr(torch.nn.functional, activation) and - callable(getattr(torch.nn.functional, activation)) + elif hasattr(torch.nn.functional, activation) and callable( + getattr(torch.nn.functional, activation) ): self.__activation = getattr(torch.nn.functional, activation) else: @@ -40,30 +45,42 @@ class _GraphSAINTAggregationLayers: if type(batch_norm) != bool: raise TypeError else: - self.__optional_batch_normalization: _typing.Optional[torch.nn.BatchNorm1d] = ( + self.__optional_batch_normalization: _typing.Optional[ + torch.nn.BatchNorm1d + ] = ( torch.nn.BatchNorm1d(output_dimension, 1e-8) - if batch_norm else None + if batch_norm + else None ) def forward( - self, x: _typing.Union[torch.Tensor, _typing.Tuple[torch.Tensor, torch.Tensor]], - _edge_index: torch.Tensor, _edge_weight: _typing.Optional[torch.Tensor] = None, - _size: _typing.Optional[_typing.Tuple[int, int]] = None + self, + x: _typing.Union[ + torch.Tensor, _typing.Tuple[torch.Tensor, torch.Tensor] + ], + _edge_index: torch.Tensor, + _edge_weight: _typing.Optional[torch.Tensor] = None, + _size: _typing.Optional[_typing.Tuple[int, int]] = None, ) -> torch.Tensor: __output: torch.Tensor = self.__linear_transform(x) if self.__activation is not None and callable(self.__activation): __output: torch.Tensor = self.__activation(__output) - if ( - self.__optional_batch_normalization is not None and - isinstance(self.__optional_batch_normalization, torch.nn.BatchNorm1d) + if self.__optional_batch_normalization is not None and isinstance( + self.__optional_batch_normalization, torch.nn.BatchNorm1d ): - __output: torch.Tensor = self.__optional_batch_normalization(__output) + __output: torch.Tensor = self.__optional_batch_normalization( + __output + ) return __output class Order1Aggregator(MessagePassing): def __init__( - self, input_dimension: int, output_dimension: int, bias: bool = True, - activation: _typing.Optional[str] = "ReLU", batch_norm: bool = True + self, + input_dimension: int, + output_dimension: int, + bias: bool = True, + activation: _typing.Optional[str] = "ReLU", + batch_norm: bool = True, ): super().__init__(aggr="add") if not type(input_dimension) == type(output_dimension) == int: @@ -72,16 +89,17 @@ class _GraphSAINTAggregationLayers: raise ValueError if not type(bias) == bool: raise TypeError - self.__linear_transform = torch.nn.Linear(input_dimension, output_dimension, bias) + self.__linear_transform = torch.nn.Linear( + input_dimension, output_dimension, bias + ) self.__linear_transform.reset_parameters() if type(activation) == str: if activation.lower() == "ReLU".lower(): self.__activation = torch.nn.functional.relu elif activation.lower() == "elu": self.__activation = torch.nn.functional.elu - elif ( - hasattr(torch.nn.functional, activation) and - callable(getattr(torch.nn.functional, activation)) + elif hasattr(torch.nn.functional, activation) and callable( + getattr(torch.nn.functional, activation) ): self.__activation = getattr(torch.nn.functional, activation) else: @@ -91,15 +109,22 @@ class _GraphSAINTAggregationLayers: if type(batch_norm) != bool: raise TypeError else: - self.__optional_batch_normalization: _typing.Optional[torch.nn.BatchNorm1d] = ( + self.__optional_batch_normalization: _typing.Optional[ + torch.nn.BatchNorm1d + ] = ( torch.nn.BatchNorm1d(output_dimension, 1e-8) - if batch_norm else None + if batch_norm + else None ) def forward( - self, x: _typing.Union[torch.Tensor, _typing.Tuple[torch.Tensor, torch.Tensor]], - _edge_index: torch.Tensor, _edge_weight: _typing.Optional[torch.Tensor] = None, - _size: _typing.Optional[_typing.Tuple[int, int]] = None + self, + x: _typing.Union[ + torch.Tensor, _typing.Tuple[torch.Tensor, torch.Tensor] + ], + _edge_index: torch.Tensor, + _edge_weight: _typing.Optional[torch.Tensor] = None, + _size: _typing.Optional[_typing.Tuple[int, int]] = None, ) -> torch.Tensor: if type(x) == torch.Tensor: @@ -111,19 +136,25 @@ class _GraphSAINTAggregationLayers: __output: torch.Tensor = self.__linear_transform(__output) if self.__activation is not None and callable(self.__activation): __output: torch.Tensor = self.__activation(__output) - if ( - self.__optional_batch_normalization is not None and - isinstance(self.__optional_batch_normalization, torch.nn.BatchNorm1d) + if self.__optional_batch_normalization is not None and isinstance( + self.__optional_batch_normalization, torch.nn.BatchNorm1d ): - __output: torch.Tensor = self.__optional_batch_normalization(__output) + __output: torch.Tensor = self.__optional_batch_normalization( + __output + ) return __output - def message(self, x_j: torch.Tensor, edge_weight: _typing.Optional[torch.Tensor]) -> torch.Tensor: + def message( + self, x_j: torch.Tensor, edge_weight: _typing.Optional[torch.Tensor] + ) -> torch.Tensor: return x_j if edge_weight is None else edge_weight.view(-1, 1) * x_j def message_and_aggregate( - self, adj_t: SparseTensor, - x: _typing.Union[torch.Tensor, _typing.Tuple[torch.Tensor, torch.Tensor]] + self, + adj_t: SparseTensor, + x: _typing.Union[ + torch.Tensor, _typing.Tuple[torch.Tensor, torch.Tensor] + ], ) -> torch.Tensor: return matmul(adj_t, x[0], reduce=self.aggr) @@ -132,14 +163,19 @@ class _GraphSAINTAggregationLayers: return (self._order + 1) * self._each_order_output_dimension def __init__( - self, _input_dimension: int, _each_order_output_dimension: int, _order: int, - bias: bool = True, activation: _typing.Optional[str] = "ReLU", - batch_norm: bool = True, _dropout: _typing.Optional[float] = ... + self, + _input_dimension: int, + _each_order_output_dimension: int, + _order: int, + bias: bool = True, + activation: _typing.Optional[str] = "ReLU", + batch_norm: bool = True, + _dropout: _typing.Optional[float] = ..., ): super().__init__() if not ( - type(_input_dimension) == type(_order) == int and - type(_each_order_output_dimension) == int + type(_input_dimension) == type(_order) == int + and type(_each_order_output_dimension) == int ): raise TypeError if _input_dimension <= 0 or _each_order_output_dimension <= 0: @@ -152,13 +188,19 @@ class _GraphSAINTAggregationLayers: if type(bias) != bool: raise TypeError self.__order0_transform = self.Order0Aggregator( - self._input_dimension, self._each_order_output_dimension, bias, - activation, batch_norm + self._input_dimension, + self._each_order_output_dimension, + bias, + activation, + batch_norm, ) if _order == 1: self.__order1_transform = self.Order1Aggregator( - self._input_dimension, self._each_order_output_dimension, bias, - activation, batch_norm + self._input_dimension, + self._each_order_output_dimension, + bias, + activation, + batch_norm, ) else: self.__order1_transform = None @@ -167,33 +209,35 @@ class _GraphSAINTAggregationLayers: _dropout = 0 if _dropout > 1: _dropout = 1 - self.__optional_dropout: _typing.Optional[torch.nn.Dropout] = ( - torch.nn.Dropout(_dropout) - ) + self.__optional_dropout: _typing.Optional[ + torch.nn.Dropout + ] = torch.nn.Dropout(_dropout) else: self.__optional_dropout: _typing.Optional[torch.nn.Dropout] = None def _forward( - self, x: _typing.Union[torch.Tensor, _typing.Tuple[torch.Tensor, torch.Tensor]], - edge_index: torch.Tensor, edge_weight: _typing.Optional[torch.Tensor] = None, - size: _typing.Optional[_typing.Tuple[int, int]] = None + self, + x: _typing.Union[torch.Tensor, _typing.Tuple[torch.Tensor, torch.Tensor]], + edge_index: torch.Tensor, + edge_weight: _typing.Optional[torch.Tensor] = None, + size: _typing.Optional[_typing.Tuple[int, int]] = None, ) -> torch.Tensor: - if ( - self.__order1_transform is not None and - isinstance(self.__order1_transform, self.Order1Aggregator) + if self.__order1_transform is not None and isinstance( + self.__order1_transform, self.Order1Aggregator ): __output: torch.Tensor = torch.cat( [ self.__order0_transform(x, edge_index, edge_weight, size), - self.__order1_transform(x, edge_index, edge_weight, size) + self.__order1_transform(x, edge_index, edge_weight, size), ], - dim=1 + dim=1, ) else: - __output: torch.Tensor = self.__order0_transform(x, edge_index, edge_weight, size) - if ( - self.__optional_dropout is not None and - isinstance(self.__optional_dropout, torch.nn.Dropout) + __output: torch.Tensor = self.__order0_transform( + x, edge_index, edge_weight, size + ) + if self.__optional_dropout is not None and isinstance( + self.__optional_dropout, torch.nn.Dropout ): __output: torch.Tensor = self.__optional_dropout(__output) return __output @@ -205,7 +249,9 @@ class _GraphSAINTAggregationLayers: edge_index: torch.LongTensor = getattr(data, "edge_index") if type(edge_index) != torch.Tensor: raise TypeError - edge_weight: _typing.Optional[torch.Tensor] = getattr(data, "edge_weight", None) + edge_weight: _typing.Optional[torch.Tensor] = getattr( + data, "edge_weight", None + ) if edge_weight is not None and type(edge_weight) != torch.Tensor: raise TypeError return self._forward(x, edge_index, edge_weight) @@ -219,8 +265,8 @@ class _GraphSAINTAggregationLayers: if type(tenser_or_data) == torch.Tensor: return self.__dropout_module(tenser_or_data) elif ( - hasattr(tenser_or_data, "x") and - type(getattr(tenser_or_data, "x")) == torch.Tensor + hasattr(tenser_or_data, "x") + and type(getattr(tenser_or_data, "x")) == torch.Tensor ): return self.__dropout_module(getattr(tenser_or_data, "x")) else: @@ -229,14 +275,17 @@ class _GraphSAINTAggregationLayers: class GraphSAINTMultiOrderAggregationModel(ClassificationSupportedSequentialModel): def __init__( - self, num_features: int, num_classes: int, - _output_dimension_for_each_order: int, - _layers_order_list: _typing.Sequence[int], - _pre_dropout: float, - _layers_dropout: _typing.Union[float, _typing.Sequence[float]], - activation: _typing.Optional[str] = "ReLU", - bias: bool = True, batch_norm: bool = True, - normalize: bool = True + self, + num_features: int, + num_classes: int, + _output_dimension_for_each_order: int, + _layers_order_list: _typing.Sequence[int], + _pre_dropout: float, + _layers_dropout: _typing.Union[float, _typing.Sequence[float]], + activation: _typing.Optional[str] = "ReLU", + bias: bool = True, + batch_norm: bool = True, + normalize: bool = True, ): super(GraphSAINTMultiOrderAggregationModel, self).__init__() if type(_output_dimension_for_each_order) != int: @@ -269,10 +318,17 @@ class GraphSAINTMultiOrderAggregationModel(ClassificationSupportedSequentialMode _pre_dropout = 1 self.__sequential_encoding_layers: torch.nn.ModuleList = torch.nn.ModuleList( ( - _GraphSAINTAggregationLayers.WrappedDropout(torch.nn.Dropout(_pre_dropout)), + _GraphSAINTAggregationLayers.WrappedDropout( + torch.nn.Dropout(_pre_dropout) + ), _GraphSAINTAggregationLayers.MultiOrderAggregationLayer( - num_features, _output_dimension_for_each_order, _layers_order_list[0], bias, - activation, batch_norm, _layers_dropout[0] + num_features, + _output_dimension_for_each_order, + _layers_order_list[0], + bias, + activation, + batch_norm, + _layers_dropout[0], ), ) ) @@ -280,14 +336,19 @@ class GraphSAINTMultiOrderAggregationModel(ClassificationSupportedSequentialMode self.__sequential_encoding_layers.append( _GraphSAINTAggregationLayers.MultiOrderAggregationLayer( self.__sequential_encoding_layers[-1].integral_output_dimension, - _output_dimension_for_each_order, _layers_order_list[_layer_index], bias, - activation, batch_norm, _layers_dropout[_layer_index] - + _output_dimension_for_each_order, + _layers_order_list[_layer_index], + bias, + activation, + batch_norm, + _layers_dropout[_layer_index], ) ) self.__apply_normalize: bool = normalize self.__linear_transform: torch.nn.Linear = torch.nn.Linear( - self.__sequential_encoding_layers[-1].integral_output_dimension, num_classes, bias + self.__sequential_encoding_layers[-1].integral_output_dimension, + num_classes, + bias, ) self.__linear_transform.reset_parameters() @@ -302,8 +363,8 @@ class GraphSAINTMultiOrderAggregationModel(ClassificationSupportedSequentialMode if type(getattr(data, "edge_index")) != torch.Tensor: raise TypeError if ( - getattr(data, "edge_weight", None) is not None and - type(getattr(data, "edge_weight")) != torch.Tensor + getattr(data, "edge_weight", None) is not None + and type(getattr(data, "edge_weight")) != torch.Tensor ): raise TypeError for encoding_layer in self.__sequential_encoding_layers: @@ -318,12 +379,12 @@ class GraphSAINTMultiOrderAggregationModel(ClassificationSupportedSequentialMode @register_model("GraphSAINTAggregationModel") class GraphSAINTAggregationModel(ClassificationModel): def __init__( - self, - num_features: int = ..., - num_classes: int = ..., - device: _typing.Union[str, torch.device] = ..., - init: bool = False, - **kwargs + self, + num_features: int = ..., + num_classes: int = ..., + device: _typing.Union[str, torch.device] = ..., + init: bool = False, + **kwargs ): super(GraphSAINTAggregationModel, self).__init__( num_features, num_classes, device=device, init=init, **kwargs @@ -333,7 +394,8 @@ class GraphSAINTAggregationModel(ClassificationModel): def _initialize(self): """ Initialize model """ self.model = GraphSAINTMultiOrderAggregationModel( - self.num_features, self.num_classes, + self.num_features, + self.num_classes, self.hyper_parameter.get("output_dimension_for_each_order"), self.hyper_parameter.get("layers_order_list"), self.hyper_parameter.get("pre_dropout"), @@ -341,5 +403,5 @@ class GraphSAINTAggregationModel(ClassificationModel): self.hyper_parameter.get("activation", "ReLU"), bool(self.hyper_parameter.get("bias", True)), bool(self.hyper_parameter.get("batch_norm", True)), - bool(self.hyper_parameter.get("normalize", True)) + bool(self.hyper_parameter.get("normalize", True)), ).to(self.device) diff --git a/autogl/module/model/graphsage.py b/autogl/module/model/graphsage.py index 727b9d0..3b63baf 100644 --- a/autogl/module/model/graphsage.py +++ b/autogl/module/model/graphsage.py @@ -14,33 +14,36 @@ LOGGER = get_logger("SAGEModel") class GraphSAGE(ClassificationSupportedSequentialModel): class _SAGELayer(torch.nn.Module): def __init__( - self, input_channels: int, output_channels: int, aggr: str, - activation_name: _typing.Optional[str] = ..., - dropout_probability: _typing.Optional[float] = ... + self, + input_channels: int, + output_channels: int, + aggr: str, + activation_name: _typing.Optional[str] = ..., + dropout_probability: _typing.Optional[float] = ..., ): super().__init__() self._convolution: SAGEConv = SAGEConv( input_channels, output_channels, aggr=aggr ) if ( - activation_name is not Ellipsis and - activation_name is not None and - type(activation_name) == str + activation_name is not Ellipsis + and activation_name is not None + and type(activation_name) == str ): self._activation_name: _typing.Optional[str] = activation_name else: self._activation_name: _typing.Optional[str] = None if ( - dropout_probability is not Ellipsis and - dropout_probability is not None and - type(dropout_probability) == float + dropout_probability is not Ellipsis + and dropout_probability is not None + and type(dropout_probability) == float ): if dropout_probability < 0: dropout_probability = 0 if dropout_probability > 1: dropout_probability = 1 - self._dropout: _typing.Optional[torch.nn.Dropout] = ( - torch.nn.Dropout(dropout_probability) + self._dropout: _typing.Optional[torch.nn.Dropout] = torch.nn.Dropout( + dropout_probability ) else: self._dropout: _typing.Optional[torch.nn.Dropout] = None @@ -59,13 +62,15 @@ class GraphSAGE(ClassificationSupportedSequentialModel): return x def __init__( - self, num_features: int, num_classes: int, - hidden_features: _typing.Sequence[int], - activation_name: str, - layers_dropout: _typing.Union[ - _typing.Optional[float], _typing.Sequence[_typing.Optional[float]] - ] = None, - aggr: str = "mean" + self, + num_features: int, + num_classes: int, + hidden_features: _typing.Sequence[int], + activation_name: str, + layers_dropout: _typing.Union[ + _typing.Optional[float], _typing.Sequence[_typing.Optional[float]] + ] = None, + aggr: str = "mean", ): super().__init__() if not type(num_features) == type(num_classes) == int: @@ -85,9 +90,9 @@ class GraphSAGE(ClassificationSupportedSequentialModel): raise TypeError _layers_dropout: _typing.Sequence[_typing.Optional[float]] = layers_dropout elif layers_dropout is None or type(layers_dropout) == float: - _layers_dropout: _typing.Sequence[_typing.Optional[float]] = ( - [layers_dropout for _ in range(len(hidden_features))] + [None] - ) + _layers_dropout: _typing.Sequence[_typing.Optional[float]] = [ + layers_dropout for _ in range(len(hidden_features)) + ] + [None] else: raise TypeError if not type(activation_name) == type(aggr) == str: @@ -96,32 +101,51 @@ class GraphSAGE(ClassificationSupportedSequentialModel): aggr = "mean" if len(hidden_features) == 0: - self.__sequential_encoding_layers: torch.nn.ModuleList = torch.nn.ModuleList([ - self._SAGELayer( - num_features, num_classes, - aggr, activation_name, _layers_dropout[0] + self.__sequential_encoding_layers: torch.nn.ModuleList = ( + torch.nn.ModuleList( + [ + self._SAGELayer( + num_features, + num_classes, + aggr, + activation_name, + _layers_dropout[0], + ) + ] ) - ]) + ) else: - self.__sequential_encoding_layers: torch.nn.ModuleList = torch.nn.ModuleList([ - self._SAGELayer( - num_features, hidden_features[0], - aggr, activation_name, _layers_dropout[0] + self.__sequential_encoding_layers: torch.nn.ModuleList = ( + torch.nn.ModuleList( + [ + self._SAGELayer( + num_features, + hidden_features[0], + aggr, + activation_name, + _layers_dropout[0], + ) + ] ) - ]) + ) for i in range(len(hidden_features)): if i + 1 < len(hidden_features): self.__sequential_encoding_layers.append( self._SAGELayer( - hidden_features[i], hidden_features[i + 1], aggr, - activation_name, _layers_dropout[i + 1] + hidden_features[i], + hidden_features[i + 1], + aggr, + activation_name, + _layers_dropout[i + 1], ) ) else: self.__sequential_encoding_layers.append( self._SAGELayer( - hidden_features[i], num_classes, aggr, - _layers_dropout[i + 1] + hidden_features[i], + num_classes, + aggr, + _layers_dropout[i + 1], ) ) @@ -131,9 +155,10 @@ class GraphSAGE(ClassificationSupportedSequentialModel): def cls_encode(self, data) -> torch.Tensor: if ( - hasattr(data, "edge_indexes") and - isinstance(getattr(data, "edge_indexes"), _typing.Sequence) and - len(getattr(data, "edge_indexes")) == len(self.__sequential_encoding_layers) + hasattr(data, "edge_indexes") + and isinstance(getattr(data, "edge_indexes"), _typing.Sequence) + and len(getattr(data, "edge_indexes")) + == len(self.__sequential_encoding_layers) ): for __edge_index in getattr(data, "edge_indexes"): if type(__edge_index) != torch.Tensor: @@ -272,9 +297,10 @@ class AutoSAGE(BaseModel): return self.initialized = True self.model = GraphSAGE( - self.num_features, self.num_classes, + self.num_features, + self.num_classes, self.hyperparams.get("hidden"), self.hyperparams.get("act", "relu"), self.hyperparams.get("dropout", None), - self.hyperparams.get("agg", "mean") + self.hyperparams.get("agg", "mean"), ).to(self.device) diff --git a/autogl/module/nas/__init__.py b/autogl/module/nas/__init__.py index be2b683..dab357f 100644 --- a/autogl/module/nas/__init__.py +++ b/autogl/module/nas/__init__.py @@ -1,8 +1,4 @@ -from . import ( - algorithm, - estimator, - space -) +from . import algorithm, estimator, space from .algorithm import NAS_ALGO_DICT from .estimator import NAS_ESTIMATOR_DICT diff --git a/autogl/module/nas/algorithm/__init__.py b/autogl/module/nas/algorithm/__init__.py index eacc45b..5056d17 100644 --- a/autogl/module/nas/algorithm/__init__.py +++ b/autogl/module/nas/algorithm/__init__.py @@ -8,10 +8,13 @@ from .base import BaseNAS NAS_ALGO_DICT = {} + def register_nas_algo(name): def register_nas_algo_cls(cls): if name in NAS_ALGO_DICT: - raise ValueError("Cannot register duplicate NAS algorithm ({})".format(name)) + raise ValueError( + "Cannot register duplicate NAS algorithm ({})".format(name) + ) if not issubclass(cls, BaseNAS): raise ValueError( "Model ({}: {}) must extend NAS algorithm".format(name, cls.__name__) @@ -21,11 +24,13 @@ def register_nas_algo(name): return register_nas_algo_cls + from .darts import Darts from .enas import Enas from .random_search import RandomSearch from .rl import RL, GraphNasRL + def build_nas_algo_from_name(name: str) -> BaseNAS: """ Parameters @@ -46,4 +51,5 @@ def build_nas_algo_from_name(name: str) -> BaseNAS: assert name in NAS_ALGO_DICT, "HPO module do not have name " + name return NAS_ALGO_DICT[name]() + __all__ = ["BaseNAS", "Darts", "Enas", "RandomSearch", "RL", "GraphNasRL"] diff --git a/autogl/module/nas/algorithm/darts.py b/autogl/module/nas/algorithm/darts.py index 924ab7f..358b1a8 100644 --- a/autogl/module/nas/algorithm/darts.py +++ b/autogl/module/nas/algorithm/darts.py @@ -17,6 +17,7 @@ from nni.retiarii.oneshot.pytorch.darts import DartsLayerChoice, DartsInputChoic _logger = logging.getLogger(__name__) + @register_nas_algo("darts") class Darts(BaseNAS): """ @@ -42,7 +43,17 @@ class Darts(BaseNAS): The device of the whole process """ - def __init__(self, num_epochs=5, workers = 4, gradient_clip = 5.0, model_lr = 1e-3, model_wd = 5e-4, arch_lr = 3e-4, arch_wd = 1e-3, device="cuda"): + def __init__( + self, + num_epochs=5, + workers=4, + gradient_clip=5.0, + model_lr=1e-3, + model_wd=5e-4, + arch_lr=3e-4, + arch_wd=1e-3, + device="cuda", + ): super().__init__(device=device) self.num_epochs = num_epochs self.workers = workers diff --git a/autogl/module/nas/algorithm/enas.py b/autogl/module/nas/algorithm/enas.py index 7e09dfb..5259e07 100644 --- a/autogl/module/nas/algorithm/enas.py +++ b/autogl/module/nas/algorithm/enas.py @@ -8,13 +8,25 @@ import torch.nn.functional as F from . import register_nas_algo from .base import BaseNAS from ..space import BaseSpace -from ..utils import AverageMeterGroup, replace_layer_choice, replace_input_choice, get_module_order, sort_replaced_module +from ..utils import ( + AverageMeterGroup, + replace_layer_choice, + replace_input_choice, + get_module_order, + sort_replaced_module, +) from tqdm import tqdm, trange -from .rl import PathSamplingLayerChoice,PathSamplingInputChoice,ReinforceField,ReinforceController +from .rl import ( + PathSamplingLayerChoice, + PathSamplingInputChoice, + ReinforceField, + ReinforceController, +) from ....utils import get_logger LOGGER = get_logger("ENAS") + @register_nas_algo("enas") class Enas(BaseNAS): """ @@ -52,29 +64,44 @@ class Enas(BaseNAS): The device of the whole process, e.g. "cuda", torch.device("cpu") """ - def __init__(self, num_epochs = 5, n_warmup = 100, log_frequency=None, grad_clip=5., entropy_weight=0.0001, skip_weight=0.8, baseline_decay=0.999, - ctrl_lr=0.00035, ctrl_steps_aggregate=20, ctrl_kwargs=None,model_lr=5e-3,model_wd=5e-4, disable_progress = True, device="cuda"): + def __init__( + self, + num_epochs=5, + n_warmup=100, + log_frequency=None, + grad_clip=5.0, + entropy_weight=0.0001, + skip_weight=0.8, + baseline_decay=0.999, + ctrl_lr=0.00035, + ctrl_steps_aggregate=20, + ctrl_kwargs=None, + model_lr=5e-3, + model_wd=5e-4, + disable_progress=True, + device="cuda", + ): super().__init__(device) - self.device=device + self.device = device self.num_epochs = num_epochs self.log_frequency = log_frequency self.entropy_weight = entropy_weight self.skip_weight = skip_weight self.baseline_decay = baseline_decay - self.baseline = 0. + self.baseline = 0.0 self.ctrl_steps_aggregate = ctrl_steps_aggregate self.grad_clip = grad_clip - self.ctrl_kwargs=ctrl_kwargs - self.ctrl_lr=ctrl_lr - self.n_warmup=n_warmup + self.ctrl_kwargs = ctrl_kwargs + self.ctrl_lr = ctrl_lr + self.n_warmup = n_warmup self.model_lr = model_lr self.model_wd = model_wd self.disable_progress = disable_progress def search(self, space: BaseSpace, dset, estimator): self.model = space - self.dataset = dset#.to(self.device) - self.estimator = estimator + self.dataset = dset # .to(self.device) + self.estimator = estimator # replace choice self.nas_modules = [] @@ -89,80 +116,98 @@ class Enas(BaseNAS): self.model.parameters(), lr=self.model_lr, weight_decay=self.model_wd ) # fields - self.nas_fields = [ReinforceField(name, len(module), - isinstance(module, PathSamplingLayerChoice) or module.n_chosen == 1) - for name, module in self.nas_modules] - self.controller = ReinforceController(self.nas_fields, **(self.ctrl_kwargs or {})) - self.ctrl_optim = torch.optim.Adam(self.controller.parameters(), lr=self.ctrl_lr) + self.nas_fields = [ + ReinforceField( + name, + len(module), + isinstance(module, PathSamplingLayerChoice) or module.n_chosen == 1, + ) + for name, module in self.nas_modules + ] + self.controller = ReinforceController( + self.nas_fields, **(self.ctrl_kwargs or {}) + ) + self.ctrl_optim = torch.optim.Adam( + self.controller.parameters(), lr=self.ctrl_lr + ) # warm up supernet with tqdm(range(self.n_warmup), disable=self.disable_progress) as bar: for i in bar: - acc,l1=self._train_model(i) + acc, l1 = self._train_model(i) with torch.no_grad(): - val_acc,val_loss=self._infer('val') - bar.set_postfix(loss=l1,acc=acc,val_acc=val_acc,val_loss=val_loss) + val_acc, val_loss = self._infer("val") + bar.set_postfix(loss=l1, acc=acc, val_acc=val_acc, val_loss=val_loss) # train with tqdm(range(self.num_epochs), disable=self.disable_progress) as bar: for i in bar: try: - l1=self._train_model(i) - l2=self._train_controller(i) + l1 = self._train_model(i) + l2 = self._train_controller(i) except Exception as e: print(e) - nm=self.nas_modules + nm = self.nas_modules for i in range(len(nm)): print(nm[i][1].sampled) - bar.set_postfix(loss_model=l1,reward_controller=l2) - - selection=self.export() - #print(selection) - return space.parse_model(selection,self.device) - - def _train_model(self, epoch): + bar.set_postfix(loss_model=l1, reward_controller=l2) + + selection = self.export() + # print(selection) + return space.parse_model(selection, self.device) + + def _train_model(self, epoch): self.model.train() self.controller.eval() self.model_optim.zero_grad() self._resample() - metric,loss=self._infer() + metric, loss = self._infer() loss.backward() if self.grad_clip > 0: nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip) self.model_optim.step() - return metric,loss.item() + return metric, loss.item() def _train_controller(self, epoch): self.model.eval() self.controller.train() self.ctrl_optim.zero_grad() - rewards=[] + rewards = [] for ctrl_step in range(self.ctrl_steps_aggregate): self._resample() with torch.no_grad(): - metric,loss=self._infer(mask='val') - reward =metric + metric, loss = self._infer(mask="val") + reward = metric rewards.append(reward) if self.entropy_weight: reward += self.entropy_weight * self.controller.sample_entropy.item() - self.baseline = self.baseline * self.baseline_decay + reward * (1 - self.baseline_decay) + self.baseline = self.baseline * self.baseline_decay + reward * ( + 1 - self.baseline_decay + ) loss = self.controller.sample_log_prob * (reward - self.baseline) if self.skip_weight: loss += self.skip_weight * self.controller.sample_skip_penalty loss /= self.ctrl_steps_aggregate loss.backward() - + if (ctrl_step + 1) % self.ctrl_steps_aggregate == 0: if self.grad_clip > 0: - nn.utils.clip_grad_norm_(self.controller.parameters(), self.grad_clip) + nn.utils.clip_grad_norm_( + self.controller.parameters(), self.grad_clip + ) self.ctrl_optim.step() self.ctrl_optim.zero_grad() if self.log_frequency is not None and ctrl_step % self.log_frequency == 0: - LOGGER.info('RL Epoch [%d/%d] Step [%d/%d] %s', epoch + 1, self.num_epochs, - ctrl_step + 1, self.ctrl_steps_aggregate) - return sum(rewards)/len(rewards) + LOGGER.info( + "RL Epoch [%d/%d] Step [%d/%d] %s", + epoch + 1, + self.num_epochs, + ctrl_step + 1, + self.ctrl_steps_aggregate, + ) + return sum(rewards) / len(rewards) def _resample(self): result = self.controller.resample() @@ -174,6 +219,6 @@ class Enas(BaseNAS): with torch.no_grad(): return self.controller.resample() - def _infer(self,mask='train'): - metric, loss = self.estimator.infer(self.model, self.dataset,mask=mask) + def _infer(self, mask="train"): + metric, loss = self.estimator.infer(self.model, self.dataset, mask=mask) return metric[0], loss diff --git a/autogl/module/nas/algorithm/random_search.py b/autogl/module/nas/algorithm/random_search.py index b3286ac..7d6bc90 100644 --- a/autogl/module/nas/algorithm/random_search.py +++ b/autogl/module/nas/algorithm/random_search.py @@ -5,17 +5,24 @@ import torch.nn.functional as F from . import register_nas_algo from .base import BaseNAS from ..space import BaseSpace -from ..utils import AverageMeterGroup, replace_layer_choice, replace_input_choice, get_module_order, sort_replaced_module +from ..utils import ( + AverageMeterGroup, + replace_layer_choice, + replace_input_choice, + get_module_order, + sort_replaced_module, +) from tqdm import tqdm -from .rl import PathSamplingLayerChoice,PathSamplingInputChoice +from .rl import PathSamplingLayerChoice, PathSamplingInputChoice import numpy as np from ....utils import get_logger LOGGER = get_logger("random_search_NAS") + @register_nas_algo("random") class RandomSearch(BaseNAS): - ''' + """ Uniformly random architecture search Parameters @@ -26,52 +33,53 @@ class RandomSearch(BaseNAS): Number of epochs planned for training. disable_progeress: boolean Control whether show the progress bar. - ''' - def __init__(self, device='cuda', num_epochs=400, disable_progress=False): + """ + + def __init__(self, device="cuda", num_epochs=400, disable_progress=False): super().__init__(device) - self.num_epochs=num_epochs - self.disable_progress=disable_progress + self.num_epochs = num_epochs + self.disable_progress = disable_progress def search(self, space: BaseSpace, dset, estimator): - self.estimator=estimator - self.dataset=dset - self.space=space - + self.estimator = estimator + self.dataset = dset + self.space = space + self.nas_modules = [] k2o = get_module_order(self.space) replace_layer_choice(self.space, PathSamplingLayerChoice, self.nas_modules) replace_input_choice(self.space, PathSamplingInputChoice, self.nas_modules) - self.nas_modules = sort_replaced_module(k2o, self.nas_modules) - selection_range={} - for k,v in self.nas_modules: - selection_range[k]=len(v) - self.selection_dict=selection_range - - #space_size=np.prod(list(selection_range.values())) + self.nas_modules = sort_replaced_module(k2o, self.nas_modules) + selection_range = {} + for k, v in self.nas_modules: + selection_range[k] = len(v) + self.selection_dict = selection_range - arch_perfs=[] - cache={} - with tqdm(range(self.num_epochs),disable=self.disable_progress) as bar: + # space_size=np.prod(list(selection_range.values())) + + arch_perfs = [] + cache = {} + with tqdm(range(self.num_epochs), disable=self.disable_progress) as bar: for i in bar: - selection=self.sample() - vec=tuple(list(selection.values())) + selection = self.sample() + vec = tuple(list(selection.values())) if vec not in cache: - self.arch=space.parse_model(selection,self.device) - metric,loss=self._infer(mask='val') - arch_perfs.append([metric,selection]) - cache[vec]=metric - bar.set_postfix(acc=metric,max_acc=max(cache.values())) - selection=arch_perfs[np.argmax([x[0] for x in arch_perfs])][1] - arch=space.parse_model(selection,self.device) - return arch - + self.arch = space.parse_model(selection, self.device) + metric, loss = self._infer(mask="val") + arch_perfs.append([metric, selection]) + cache[vec] = metric + bar.set_postfix(acc=metric, max_acc=max(cache.values())) + selection = arch_perfs[np.argmax([x[0] for x in arch_perfs])][1] + arch = space.parse_model(selection, self.device) + return arch + def sample(self): # uniformly sample - selection={} - for k,v in self.selection_dict.items(): - selection[k]=np.random.choice(range(v)) + selection = {} + for k, v in self.selection_dict.items(): + selection[k] = np.random.choice(range(v)) return selection - def _infer(self,mask='train'): + def _infer(self, mask="train"): metric, loss = self.estimator.infer(self.arch._model, self.dataset, mask=mask) return metric[0], loss diff --git a/autogl/module/nas/algorithm/rl.py b/autogl/module/nas/algorithm/rl.py index a8fdbcf..2f48ba4 100644 --- a/autogl/module/nas/algorithm/rl.py +++ b/autogl/module/nas/algorithm/rl.py @@ -6,7 +6,13 @@ import torch.nn.functional as F from . import register_nas_algo from .base import BaseNAS from ..space import BaseSpace -from ..utils import AverageMeterGroup, replace_layer_choice, replace_input_choice, get_module_order, sort_replaced_module +from ..utils import ( + AverageMeterGroup, + replace_layer_choice, + replace_input_choice, + get_module_order, + sort_replaced_module, +) from nni.nas.pytorch.fixed import apply_fixed_architecture from tqdm import tqdm from datetime import datetime @@ -14,10 +20,16 @@ import numpy as np from ....utils import get_logger LOGGER = get_logger("random_search_NAS") + + def _get_mask(sampled, total): - multihot = [i == sampled or (isinstance(sampled, list) and i in sampled) for i in range(total)] + multihot = [ + i == sampled or (isinstance(sampled, list) and i in sampled) + for i in range(total) + ] return torch.tensor(multihot, dtype=torch.bool) # pylint: disable=not-callable + class PathSamplingLayerChoice(nn.Module): """ Mixed module, in which fprop is decided by exactly one or multiple (sampled) module. @@ -37,15 +49,21 @@ class PathSamplingLayerChoice(nn.Module): for name, module in layer_choice.named_children(): self.add_module(name, module) self.op_names.append(name) - assert self.op_names, 'There has to be at least one op to choose from.' + assert self.op_names, "There has to be at least one op to choose from." self.sampled = None # sampled can be either a list of indices or an index def forward(self, *args, **kwargs): - assert self.sampled is not None, 'At least one path needs to be sampled before fprop.' + assert ( + self.sampled is not None + ), "At least one path needs to be sampled before fprop." if isinstance(self.sampled, list): - return sum([getattr(self, self.op_names[i])(*args, **kwargs) for i in self.sampled]) # pylint: disable=not-an-iterable + return sum( + [getattr(self, self.op_names[i])(*args, **kwargs) for i in self.sampled] + ) # pylint: disable=not-an-iterable else: - return getattr(self, self.op_names[self.sampled])(*args, **kwargs) # pylint: disable=invalid-sequence-index + return getattr(self, self.op_names[self.sampled])( + *args, **kwargs + ) # pylint: disable=invalid-sequence-index def __len__(self): return len(self.op_names) @@ -75,7 +93,9 @@ class PathSamplingInputChoice(nn.Module): def forward(self, input_tensors): if isinstance(self.sampled, list): - return sum([input_tensors[t] for t in self.sampled]) # pylint: disable=not-an-iterable + return sum( + [input_tensors[t] for t in self.sampled] + ) # pylint: disable=not-an-iterable else: return input_tensors[self.sampled] @@ -87,14 +107,16 @@ class PathSamplingInputChoice(nn.Module): return _get_mask(self.sampled, len(self)) def __repr__(self): - return f'PathSamplingInputChoice(n_candidates={self.n_candidates}, chosen={self.sampled})' + return f"PathSamplingInputChoice(n_candidates={self.n_candidates}, chosen={self.sampled})" + class StackedLSTMCell(nn.Module): def __init__(self, layers, size, bias): super().__init__() self.lstm_num_layers = layers - self.lstm_modules = nn.ModuleList([nn.LSTMCell(size, size, bias=bias) - for _ in range(self.lstm_num_layers)]) + self.lstm_modules = nn.ModuleList( + [nn.LSTMCell(size, size, bias=bias) for _ in range(self.lstm_num_layers)] + ) def forward(self, inputs, hidden): prev_h, prev_c = hidden @@ -108,6 +130,7 @@ class StackedLSTMCell(nn.Module): inputs = curr_h[-1].view(1, -1) return next_h, next_c + class ReinforceField: """ A field with ``name``, with ``total`` choices. ``choose_one`` is true if one and only one is meant to be @@ -120,7 +143,8 @@ class ReinforceField: self.choose_one = choose_one def __repr__(self): - return f'ReinforceField(name={self.name}, total={self.total}, choose_one={self.choose_one})' + return f"ReinforceField(name={self.name}, total={self.total}, choose_one={self.choose_one})" + class ReinforceController(nn.Module): """ @@ -144,8 +168,16 @@ class ReinforceController(nn.Module): Can be one of ``sum`` and ``mean``. How the entropy of multi-input-choice is reduced. """ - def __init__(self, fields, lstm_size=64, lstm_num_layers=1, tanh_constant=1.5, - skip_target=0.4, temperature=None, entropy_reduction='sum'): + def __init__( + self, + fields, + lstm_size=64, + lstm_num_layers=1, + tanh_constant=1.5, + skip_target=0.4, + temperature=None, + entropy_reduction="sum", + ): super(ReinforceController, self).__init__() self.fields = fields self.lstm_size = lstm_size @@ -159,17 +191,27 @@ class ReinforceController(nn.Module): self.attn_query = nn.Linear(self.lstm_size, self.lstm_size, bias=False) self.v_attn = nn.Linear(self.lstm_size, 1, bias=False) self.g_emb = nn.Parameter(torch.randn(1, self.lstm_size) * 0.1) - self.skip_targets = nn.Parameter(torch.tensor([1.0 - self.skip_target, self.skip_target]), # pylint: disable=not-callable - requires_grad=False) - assert entropy_reduction in ['sum', 'mean'], 'Entropy reduction must be one of sum and mean.' - self.entropy_reduction = torch.sum if entropy_reduction == 'sum' else torch.mean - self.cross_entropy_loss = nn.CrossEntropyLoss(reduction='none') - self.soft = nn.ModuleDict({ - field.name: nn.Linear(self.lstm_size, field.total, bias=False) for field in fields - }) - self.embedding = nn.ModuleDict({ - field.name: nn.Embedding(field.total, self.lstm_size) for field in fields - }) + self.skip_targets = nn.Parameter( + torch.tensor( + [1.0 - self.skip_target, self.skip_target] + ), # pylint: disable=not-callable + requires_grad=False, + ) + assert entropy_reduction in [ + "sum", + "mean", + ], "Entropy reduction must be one of sum and mean." + self.entropy_reduction = torch.sum if entropy_reduction == "sum" else torch.mean + self.cross_entropy_loss = nn.CrossEntropyLoss(reduction="none") + self.soft = nn.ModuleDict( + { + field.name: nn.Linear(self.lstm_size, field.total, bias=False) + for field in fields + } + ) + self.embedding = nn.ModuleDict( + {field.name: nn.Embedding(field.total, self.lstm_size) for field in fields} + ) def resample(self): self._initialize() @@ -180,12 +222,22 @@ class ReinforceController(nn.Module): def _initialize(self): self._inputs = self.g_emb.data - self._c = [torch.zeros((1, self.lstm_size), - dtype=self._inputs.dtype, - device=self._inputs.device) for _ in range(self.lstm_num_layers)] - self._h = [torch.zeros((1, self.lstm_size), - dtype=self._inputs.dtype, - device=self._inputs.device) for _ in range(self.lstm_num_layers)] + self._c = [ + torch.zeros( + (1, self.lstm_size), + dtype=self._inputs.dtype, + device=self._inputs.device, + ) + for _ in range(self.lstm_num_layers) + ] + self._h = [ + torch.zeros( + (1, self.lstm_size), + dtype=self._inputs.dtype, + device=self._inputs.device, + ) + for _ in range(self.lstm_num_layers) + ] self.sample_log_prob = 0 self.sample_entropy = 0 self.sample_skip_penalty = 0 @@ -206,7 +258,9 @@ class ReinforceController(nn.Module): self._inputs = self.embedding[field.name](sampled) else: logit = logit.view(-1, 1) - logit = torch.cat([-logit, logit], 1) # pylint: disable=invalid-unary-operand-type + logit = torch.cat( + [-logit, logit], 1 + ) # pylint: disable=invalid-unary-operand-type sampled = torch.multinomial(F.softmax(logit, dim=-1), 1).view(-1) skip_prob = torch.sigmoid(logit) kl = torch.sum(skip_prob * torch.log(skip_prob / self.skip_targets)) @@ -214,18 +268,26 @@ class ReinforceController(nn.Module): log_prob = self.cross_entropy_loss(logit, sampled) sampled = sampled.nonzero().view(-1) if sampled.sum().item(): - self._inputs = (torch.sum(self.embedding[field.name](sampled.view(-1)), 0) / (1. + torch.sum(sampled))).unsqueeze(0) + self._inputs = ( + torch.sum(self.embedding[field.name](sampled.view(-1)), 0) + / (1.0 + torch.sum(sampled)) + ).unsqueeze(0) else: - self._inputs = torch.zeros(1, self.lstm_size, device=self.embedding[field.name].weight.device) + self._inputs = torch.zeros( + 1, self.lstm_size, device=self.embedding[field.name].weight.device + ) sampled = sampled.detach().numpy().tolist() self.sample_log_prob += self.entropy_reduction(log_prob) - entropy = (log_prob * torch.exp(-log_prob)).detach() # pylint: disable=invalid-unary-operand-type + entropy = ( + log_prob * torch.exp(-log_prob) + ).detach() # pylint: disable=invalid-unary-operand-type self.sample_entropy += self.entropy_reduction(entropy) if len(sampled) == 1: sampled = sampled[0] return sampled + @register_nas_algo("rl") class RL(BaseNAS): """ @@ -265,30 +327,44 @@ class RL(BaseNAS): Control whether show the progress bar. """ - def __init__(self, num_epochs = 5, device='cuda', log_frequency=None, - grad_clip=5., entropy_weight=0.0001, skip_weight=0.8, baseline_decay=0.999, - ctrl_lr=0.00035, ctrl_steps_aggregate=20, ctrl_kwargs=None,n_warmup=100,model_lr=5e-3,model_wd=5e-4, disable_progress=True): + def __init__( + self, + num_epochs=5, + device="cuda", + log_frequency=None, + grad_clip=5.0, + entropy_weight=0.0001, + skip_weight=0.8, + baseline_decay=0.999, + ctrl_lr=0.00035, + ctrl_steps_aggregate=20, + ctrl_kwargs=None, + n_warmup=100, + model_lr=5e-3, + model_wd=5e-4, + disable_progress=True, + ): super().__init__(device) - self.device=device + self.device = device self.num_epochs = num_epochs self.log_frequency = log_frequency self.entropy_weight = entropy_weight self.skip_weight = skip_weight self.baseline_decay = baseline_decay - self.baseline = 0. + self.baseline = 0.0 self.ctrl_steps_aggregate = ctrl_steps_aggregate self.grad_clip = grad_clip - self.ctrl_kwargs=ctrl_kwargs - self.ctrl_lr=ctrl_lr - self.n_warmup=n_warmup + self.ctrl_kwargs = ctrl_kwargs + self.ctrl_lr = ctrl_lr + self.n_warmup = n_warmup self.model_lr = model_lr self.model_wd = model_wd - self.disable_progress=disable_progress + self.disable_progress = disable_progress def search(self, space: BaseSpace, dset, estimator): self.model = space - self.dataset = dset#.to(self.device) - self.estimator = estimator + self.dataset = dset # .to(self.device) + self.estimator = estimator # replace choice self.nas_modules = [] @@ -300,69 +376,95 @@ class RL(BaseNAS): # to device self.model = self.model.to(self.device) # fields - self.nas_fields = [ReinforceField(name, len(module), - isinstance(module, PathSamplingLayerChoice) or module.n_chosen == 1) - for name, module in self.nas_modules] - self.controller = ReinforceController(self.nas_fields, **(self.ctrl_kwargs or {})) - self.ctrl_optim = torch.optim.Adam(self.controller.parameters(), lr=self.ctrl_lr) + self.nas_fields = [ + ReinforceField( + name, + len(module), + isinstance(module, PathSamplingLayerChoice) or module.n_chosen == 1, + ) + for name, module in self.nas_modules + ] + self.controller = ReinforceController( + self.nas_fields, **(self.ctrl_kwargs or {}) + ) + self.ctrl_optim = torch.optim.Adam( + self.controller.parameters(), lr=self.ctrl_lr + ) # train - with tqdm(range(self.num_epochs), disable = self.disable_progress) as bar: + with tqdm(range(self.num_epochs), disable=self.disable_progress) as bar: for i in bar: - l2=self._train_controller(i) + l2 = self._train_controller(i) bar.set_postfix(reward_controller=l2) - - selection=self.export() - arch=space.parse_model(selection,self.device) - #print(selection,arch) + + selection = self.export() + arch = space.parse_model(selection, self.device) + # print(selection,arch) return arch - + def _train_controller(self, epoch): self.model.eval() self.controller.train() self.ctrl_optim.zero_grad() - rewards=[] - with tqdm(range(self.ctrl_steps_aggregate), disable=self.disable_progress) as bar: + rewards = [] + with tqdm( + range(self.ctrl_steps_aggregate), disable=self.disable_progress + ) as bar: for ctrl_step in bar: self._resample() - metric,loss=self._infer(mask='val') - bar.set_postfix(acc=metric,loss=loss.item()) - LOGGER.info(f'{self.arch}\n{self.selection}\n{metric},{loss}') - reward =metric + metric, loss = self._infer(mask="val") + bar.set_postfix(acc=metric, loss=loss.item()) + LOGGER.info(f"{self.arch}\n{self.selection}\n{metric},{loss}") + reward = metric rewards.append(reward) if self.entropy_weight: - reward += self.entropy_weight * self.controller.sample_entropy.item() - self.baseline = self.baseline * self.baseline_decay + reward * (1 - self.baseline_decay) + reward += ( + self.entropy_weight * self.controller.sample_entropy.item() + ) + self.baseline = self.baseline * self.baseline_decay + reward * ( + 1 - self.baseline_decay + ) loss = self.controller.sample_log_prob * (reward - self.baseline) if self.skip_weight: loss += self.skip_weight * self.controller.sample_skip_penalty loss /= self.ctrl_steps_aggregate loss.backward() - + if (ctrl_step + 1) % self.ctrl_steps_aggregate == 0: if self.grad_clip > 0: - nn.utils.clip_grad_norm_(self.controller.parameters(), self.grad_clip) + nn.utils.clip_grad_norm_( + self.controller.parameters(), self.grad_clip + ) self.ctrl_optim.step() self.ctrl_optim.zero_grad() - if self.log_frequency is not None and ctrl_step % self.log_frequency == 0: - LOGGER.info('RL Epoch [%d/%d] Step [%d/%d] %s', epoch + 1, self.num_epochs, - ctrl_step + 1, self.ctrl_steps_aggregate) - return sum(rewards)/len(rewards) + if ( + self.log_frequency is not None + and ctrl_step % self.log_frequency == 0 + ): + LOGGER.info( + "RL Epoch [%d/%d] Step [%d/%d] %s", + epoch + 1, + self.num_epochs, + ctrl_step + 1, + self.ctrl_steps_aggregate, + ) + return sum(rewards) / len(rewards) def _resample(self): result = self.controller.resample() - self.arch=self.model.parse_model(result,device=self.device) - self.selection=result + self.arch = self.model.parse_model(result, device=self.device) + self.selection = result def export(self): self.controller.eval() with torch.no_grad(): return self.controller.resample() - def _infer(self,mask='train'): - metric, loss = self.estimator.infer(self.arch._model, self.dataset,mask=mask) + def _infer(self, mask="train"): + metric, loss = self.estimator.infer(self.arch._model, self.dataset, mask=mask) return metric[0], loss + @register_nas_algo("graphnas") class GraphNasRL(BaseNAS): """ @@ -404,11 +506,26 @@ class GraphNasRL(BaseNAS): Control whether show the progress bar. """ - def __init__(self, device='cuda', num_epochs=10, log_frequency=None, - grad_clip=5., entropy_weight=0.0001, skip_weight=0, baseline_decay=0.95, - ctrl_lr=0.00035, ctrl_steps_aggregate=100, ctrl_kwargs=None, n_warmup=100, model_lr=5e-3, model_wd=5e-4, topk=5, disable_progress = True): + def __init__( + self, + device="cuda", + num_epochs=10, + log_frequency=None, + grad_clip=5.0, + entropy_weight=0.0001, + skip_weight=0, + baseline_decay=0.95, + ctrl_lr=0.00035, + ctrl_steps_aggregate=100, + ctrl_kwargs=None, + n_warmup=100, + model_lr=5e-3, + model_wd=5e-4, + topk=5, + disable_progress=True, + ): super().__init__(device) - self.device=device + self.device = device self.num_epochs = num_epochs self.log_frequency = log_frequency self.entropy_weight = entropy_weight @@ -416,19 +533,19 @@ class GraphNasRL(BaseNAS): self.baseline_decay = baseline_decay self.ctrl_steps_aggregate = ctrl_steps_aggregate self.grad_clip = grad_clip - self.ctrl_kwargs=ctrl_kwargs - self.ctrl_lr=ctrl_lr - self.n_warmup=n_warmup + self.ctrl_kwargs = ctrl_kwargs + self.ctrl_lr = ctrl_lr + self.n_warmup = n_warmup self.model_lr = model_lr self.model_wd = model_wd - self.hist=[] - self.topk=topk - self.disable_progress=disable_progress + self.hist = [] + self.topk = topk + self.disable_progress = disable_progress def search(self, space: BaseSpace, dset, estimator): self.model = space - self.dataset = dset#.to(self.device) - self.estimator = estimator + self.dataset = dset # .to(self.device) + self.estimator = estimator # replace choice self.nas_modules = [] @@ -440,91 +557,114 @@ class GraphNasRL(BaseNAS): # to device self.model = self.model.to(self.device) # fields - self.nas_fields = [ReinforceField(name, len(module), - isinstance(module, PathSamplingLayerChoice) or module.n_chosen == 1) - for name, module in self.nas_modules] - self.controller = ReinforceController(self.nas_fields,lstm_size=100,temperature=5.0,tanh_constant=2.5, **(self.ctrl_kwargs or {})) - self.ctrl_optim = torch.optim.Adam(self.controller.parameters(), lr=self.ctrl_lr) + self.nas_fields = [ + ReinforceField( + name, + len(module), + isinstance(module, PathSamplingLayerChoice) or module.n_chosen == 1, + ) + for name, module in self.nas_modules + ] + self.controller = ReinforceController( + self.nas_fields, + lstm_size=100, + temperature=5.0, + tanh_constant=2.5, + **(self.ctrl_kwargs or {}), + ) + self.ctrl_optim = torch.optim.Adam( + self.controller.parameters(), lr=self.ctrl_lr + ) # train with tqdm(range(self.num_epochs), disable=self.disable_progress) as bar: for i in bar: - l2=self._train_controller(i) + l2 = self._train_controller(i) bar.set_postfix(reward_controller=l2) - + # selection=self.export() - - selections=[x[1] for x in self.hist] - candidiate_accs=[-x[0] for x in self.hist] - #print('candidiate accuracies',candidiate_accs) - selection=self._choose_best(selections) - arch=space.parse_model(selection,self.device) - #print(selection,arch) + + selections = [x[1] for x in self.hist] + candidiate_accs = [-x[0] for x in self.hist] + # print('candidiate accuracies',candidiate_accs) + selection = self._choose_best(selections) + arch = space.parse_model(selection, self.device) + # print(selection,arch) return arch - def _choose_best(self,selections): + def _choose_best(self, selections): # graphnas use top 5 models, can evaluate 20 times epoch and choose the best. - results=[] + results = [] for selection in selections: - accs=[] + accs = [] for i in tqdm(range(20), disable=self.disable_progress): - self.arch=self.model.parse_model(selection,device=self.device) - metric,loss=self._infer(mask='val') + self.arch = self.model.parse_model(selection, device=self.device) + metric, loss = self._infer(mask="val") accs.append(metric) - result=np.mean(accs) - LOGGER.info('selection {} \n acc {:.4f} +- {:.4f}'.format(selection,np.mean(accs),np.std(accs)/np.sqrt(20))) + result = np.mean(accs) + LOGGER.info( + "selection {} \n acc {:.4f} +- {:.4f}".format( + selection, np.mean(accs), np.std(accs) / np.sqrt(20) + ) + ) results.append(result) - best_selection=selections[np.argmax(results)] + best_selection = selections[np.argmax(results)] return best_selection - + def _train_controller(self, epoch): self.model.eval() self.controller.train() self.ctrl_optim.zero_grad() - rewards=[] - baseline=None + rewards = [] + baseline = None # diff: graph nas train 100 and derive 100 for every epoch(10 epochs), we just train 100(20 epochs). totol num of samples are same (2000) - with tqdm(range(self.ctrl_steps_aggregate), disable=self.disable_progress) as bar: + with tqdm( + range(self.ctrl_steps_aggregate), disable=self.disable_progress + ) as bar: for ctrl_step in bar: self._resample() - metric,loss=self._infer(mask='val') + metric, loss = self._infer(mask="val") # bar.set_postfix(acc=metric,loss=loss.item()) - LOGGER.debug(f'{self.arch}\n{self.selection}\n{metric},{loss}') + LOGGER.debug(f"{self.arch}\n{self.selection}\n{metric},{loss}") # diff: not do reward shaping as in graphnas code - reward =metric - self.hist.append([-metric,self.selection]) - if len(self.hist)>self.topk: - self.hist.sort(key=lambda x:x[0]) + reward = metric + self.hist.append([-metric, self.selection]) + if len(self.hist) > self.topk: + self.hist.sort(key=lambda x: x[0]) self.hist.pop() rewards.append(reward) - + if self.entropy_weight: - reward += self.entropy_weight * self.controller.sample_entropy.item() + reward += ( + self.entropy_weight * self.controller.sample_entropy.item() + ) if not baseline: - baseline= reward + baseline = reward else: - baseline = baseline * self.baseline_decay + reward * (1 - self.baseline_decay) + baseline = baseline * self.baseline_decay + reward * ( + 1 - self.baseline_decay + ) loss = self.controller.sample_log_prob * (reward - baseline) self.ctrl_optim.zero_grad() loss.backward() - + self.ctrl_optim.step() - - bar.set_postfix(acc=metric,max_acc=max(rewards)) - return sum(rewards)/len(rewards) + + bar.set_postfix(acc=metric, max_acc=max(rewards)) + return sum(rewards) / len(rewards) def _resample(self): result = self.controller.resample() - self.arch=self.model.parse_model(result,device=self.device) - self.selection=result + self.arch = self.model.parse_model(result, device=self.device) + self.selection = result def export(self): self.controller.eval() with torch.no_grad(): return self.controller.resample() - def _infer(self,mask='train'): - metric, loss = self.estimator.infer(self.arch._model, self.dataset,mask=mask) - return metric[0], loss \ No newline at end of file + def _infer(self, mask="train"): + metric, loss = self.estimator.infer(self.arch._model, self.dataset, mask=mask) + return metric[0], loss diff --git a/autogl/module/nas/estimator/__init__.py b/autogl/module/nas/estimator/__init__.py index e768aa2..88b9115 100644 --- a/autogl/module/nas/estimator/__init__.py +++ b/autogl/module/nas/estimator/__init__.py @@ -4,10 +4,13 @@ from .base import BaseEstimator NAS_ESTIMATOR_DICT = {} + def register_nas_estimator(name): def register_nas_estimator_cls(cls): if name in NAS_ESTIMATOR_DICT: - raise ValueError("Cannot register duplicate NAS estimator ({})".format(name)) + raise ValueError( + "Cannot register duplicate NAS estimator ({})".format(name) + ) if not issubclass(cls, BaseEstimator): raise ValueError( "Model ({}: {}) must extend NAS estimator".format(name, cls.__name__) @@ -17,9 +20,11 @@ def register_nas_estimator(name): return register_nas_estimator_cls + from .one_shot import OneShotEstimator from .train_scratch import TrainEstimator + def build_nas_estimator_from_name(name: str) -> BaseEstimator: """ Parameters @@ -40,4 +45,5 @@ def build_nas_estimator_from_name(name: str) -> BaseEstimator: assert name in NAS_ESTIMATOR_DICT, "HPO module do not have name " + name return NAS_ESTIMATOR_DICT[name]() + __all__ = ["BaseEstimator", "OneShotEstimator", "TrainEstimator"] diff --git a/autogl/module/nas/estimator/base.py b/autogl/module/nas/estimator/base.py index b846dc1..0505253 100644 --- a/autogl/module/nas/estimator/base.py +++ b/autogl/module/nas/estimator/base.py @@ -9,6 +9,7 @@ from ...train.evaluation import Evaluation, Acc import torch.nn.functional as F import torch + class BaseEstimator: """ The estimator of NAS model. @@ -21,13 +22,14 @@ class BaseEstimator: evaluation: list of autogl.module.train.evaluation.Evaluation Default evaluation metric """ - def __init__(self, loss_f: str = 'nll_loss', evaluation = [Acc()]): + + def __init__(self, loss_f: str = "nll_loss", evaluation=[Acc()]): self.loss_f = loss_f self.evaluation = evaluation def setLossFunction(self, loss_f: str): self.loss_f = loss_f - + def setEvaluation(self, evaluation): self.evaluation = evaluation diff --git a/autogl/module/nas/estimator/one_shot.py b/autogl/module/nas/estimator/one_shot.py index 5a49dca..f719e4f 100644 --- a/autogl/module/nas/estimator/one_shot.py +++ b/autogl/module/nas/estimator/one_shot.py @@ -5,6 +5,7 @@ from . import register_nas_estimator from ..space import BaseSpace from .base import BaseEstimator + @register_nas_estimator("oneshot") class OneShotEstimator(BaseEstimator): """ @@ -17,10 +18,10 @@ class OneShotEstimator(BaseEstimator): device = next(model.parameters()).device dset = dataset[0].to(device) pred = model(dset)[getattr(dset, f"{mask}_mask")] - y = dset.y[getattr(dset, f'{mask}_mask')] + y = dset.y[getattr(dset, f"{mask}_mask")] loss = getattr(F, self.loss_f)(pred, y) - #acc=sum(pred.max(1)[1]==y).item()/y.size(0) - probs = F.softmax(pred, dim = 1).detach().cpu().numpy() + # acc=sum(pred.max(1)[1]==y).item()/y.size(0) + probs = F.softmax(pred, dim=1).detach().cpu().numpy() y = y.cpu() metrics = [eva.evaluate(probs, y) for eva in self.evaluation] return metrics, loss diff --git a/autogl/module/nas/estimator/train_scratch.py b/autogl/module/nas/estimator/train_scratch.py index 2077b7a..38036cc 100644 --- a/autogl/module/nas/estimator/train_scratch.py +++ b/autogl/module/nas/estimator/train_scratch.py @@ -9,32 +9,35 @@ import torch from autogl.module.train import NodeClassificationFullTrainer, Acc + @register_nas_estimator("scratch") class TrainEstimator(BaseEstimator): """ An estimator which trans from scratch """ - def __init__(self, loss_f = "nll_loss", evaluation = [Acc()]): + + def __init__(self, loss_f="nll_loss", evaluation=[Acc()]): super().__init__(loss_f, evaluation) self.evaluation = evaluation - self.estimator=OneShotEstimator(self.loss_f, self.evaluation) + self.estimator = OneShotEstimator(self.loss_f, self.evaluation) def infer(self, model: BaseSpace, dataset, mask="train"): # self.trainer.model=model # self.trainer.device=model.device boxmodel = model.wrap() - self.trainer=NodeClassificationFullTrainer( - model=boxmodel, - optimizer=torch.optim.Adam, - lr=0.005, - max_epoch=300, - early_stopping_round=30, - weight_decay=5e-4, - device="auto", - init=False, - feval=self.evaluation, - loss=self.loss_f, - lr_scheduler_type=None) + self.trainer = NodeClassificationFullTrainer( + model=boxmodel, + optimizer=torch.optim.Adam, + lr=0.005, + max_epoch=300, + early_stopping_round=30, + weight_decay=5e-4, + device="auto", + init=False, + feval=self.evaluation, + loss=self.loss_f, + lr_scheduler_type=None, + ) try: self.trainer.train(dataset) with torch.no_grad(): @@ -42,7 +45,7 @@ class TrainEstimator(BaseEstimator): except RuntimeError as e: if "cuda" in str(e) or "CUDA" in str(e): INF = 100 - fin = [-INF if eva.is_higher_better else INF for eva in self.evaluation] + fin = [-INF if eva.is_higher_better else INF for eva in self.evaluation] return fin, 0 else: raise e diff --git a/autogl/module/nas/space/__init__.py b/autogl/module/nas/space/__init__.py index 60c1054..6d1e8f1 100644 --- a/autogl/module/nas/space/__init__.py +++ b/autogl/module/nas/space/__init__.py @@ -4,6 +4,7 @@ from .base import BaseSpace NAS_SPACE_DICT = {} + def register_nas_space(name): def register_nas_space_cls(cls): if name in NAS_SPACE_DICT: @@ -17,10 +18,12 @@ def register_nas_space(name): return register_nas_space_cls + from .graph_nas_macro import GraphNasMacroNodeClassificationSpace from .graph_nas import GraphNasNodeClassificationSpace from .single_path import SinglePathNodeClassificationSpace + def build_nas_space_from_name(name: str) -> BaseSpace: """ Parameters @@ -41,4 +44,10 @@ def build_nas_space_from_name(name: str) -> BaseSpace: assert name in NAS_SPACE_DICT, "HPO module do not have name " + name return NAS_SPACE_DICT[name]() -__all__ = ["BaseSpace", "GraphNasMacroNodeClassificationSpace", "GraphNasNodeClassificationSpace", "SinglePathNodeClassificationSpace"] + +__all__ = [ + "BaseSpace", + "GraphNasMacroNodeClassificationSpace", + "GraphNasNodeClassificationSpace", + "SinglePathNodeClassificationSpace", +] diff --git a/autogl/module/nas/space/base.py b/autogl/module/nas/space/base.py index 433157d..02d4ca6 100644 --- a/autogl/module/nas/space/base.py +++ b/autogl/module/nas/space/base.py @@ -11,7 +11,8 @@ from ....utils import get_logger from ...model import AutoGCN -class OrderedMutable(): + +class OrderedMutable: """ An abstract class with order, enabling to sort mutables with a certain rank. @@ -20,20 +21,35 @@ class OrderedMutable(): order : int The order of the mutable """ + def __init__(self, order): self.order = order + class OrderedLayerChoice(OrderedMutable, mutables.LayerChoice): - def __init__(self, order, op_candidates, reduction="sum", return_mask=False, key=None): + def __init__( + self, order, op_candidates, reduction="sum", return_mask=False, key=None + ): OrderedMutable.__init__(self, order) mutables.LayerChoice.__init__(self, op_candidates, reduction, return_mask, key) + class OrderedInputChoice(OrderedMutable, mutables.InputChoice): - def __init__(self, order, n_candidates=None, choose_from=None, n_chosen=None, - reduction="sum", return_mask=False, key=None): + def __init__( + self, + order, + n_candidates=None, + choose_from=None, + n_chosen=None, + reduction="sum", + return_mask=False, + key=None, + ): OrderedMutable.__init__(self, order) - mutables.InputChoice.__init__(self, n_candidates, choose_from, n_chosen, - reduction, return_mask, key) + mutables.InputChoice.__init__( + self, n_candidates, choose_from, n_chosen, reduction, return_mask, key + ) + class StrModule(nn.Module): """ @@ -45,15 +61,17 @@ class StrModule(nn.Module): name : anything the name of module, can be any type """ + def __init__(self, name): super().__init__() self.str = name - def forward(self, *args,**kwargs): - return self.str + def forward(self, *args, **kwargs): + return self.str def __repr__(self): - return '{}({})'.format(self.__class__.__name__,self.str) + return "{}({})".format(self.__class__.__name__, self.str) + def map_nn(names): """ @@ -66,6 +84,7 @@ def map_nn(names): """ return [StrModule(x) for x in names] + class BoxModel(BaseModel): """ The box wrapping a space, can be passed to later procedure or trainer @@ -77,6 +96,7 @@ class BoxModel(BaseModel): device : str or torch.device The device to place the model """ + _logger = get_logger("space model") def __init__(self, space_model, device=torch.device("cuda")): @@ -93,7 +113,7 @@ class BoxModel(BaseModel): def fix(self, selection): """ - To fix self._model with a selection + To fix self._model with a selection Parameters ---------- @@ -129,6 +149,7 @@ class BoxModel(BaseModel): def model(self): return self._model + class BaseSpace(nn.Module): """ Base space class of NAS module. Defining space containing all models. @@ -187,7 +208,9 @@ class BaseSpace(nn.Module): if not self._initialized: self._initialized = True - def setLayerChoice(self, order, op_candidates, reduction="sum", return_mask=False, key=None): + def setLayerChoice( + self, order, op_candidates, reduction="sum", return_mask=False, key=None + ): """ Give a unique key if not given """ @@ -199,8 +222,16 @@ class BaseSpace(nn.Module): layer = OrderedLayerChoice(order, op_candidates, reduction, return_mask, orikey) return layer - def setInputChoice(self, order, n_candidates=None, choose_from=None, n_chosen=None, - reduction="sum", return_mask=False, key=None): + def setInputChoice( + self, + order, + n_candidates=None, + choose_from=None, + n_chosen=None, + reduction="sum", + return_mask=False, + key=None, + ): """ Give a unique key if not given """ @@ -209,8 +240,9 @@ class BaseSpace(nn.Module): key = f"default_key_{self._default_key}" self._default_key += 1 orikey = key - layer = OrderedInputChoice(order, n_candidates, choose_from, n_chosen, - reduction, return_mask, orikey) + layer = OrderedInputChoice( + order, n_candidates, choose_from, n_chosen, reduction, return_mask, orikey + ) return layer def wrap(self, device="cuda"): @@ -218,8 +250,9 @@ class BaseSpace(nn.Module): Return a BoxModel which wrap self as a model Used to pass to trainer To use this function, must contain `input_dim` and `output_dim` - """ - return BoxModel(self, device) + """ + return BoxModel(self, device) + class FixedInputChoice(nn.Module): """ @@ -230,6 +263,7 @@ class FixedInputChoice(nn.Module): mask : list The mask indicating which input to choose """ + def __init__(self, mask): self.mask_len = len(mask) for i in range(self.mask_len): @@ -242,6 +276,7 @@ class FixedInputChoice(nn.Module): if len(optional_inputs) == self.mask_len: return optional_inputs[self.selected] + class CleanFixedArchitecture(FixedArchitecture): """ Fixed architecture mutator that always selects a certain graph, allowing deepcopy @@ -295,6 +330,7 @@ class CleanFixedArchitecture(FixedArchitecture): else: self.replace_all_choice(mutable, global_name) + def apply_fixed_architecture(model, fixed_arc, verbose=True): """ Load architecture from `fixed_arc` and apply to model. diff --git a/autogl/module/nas/space/graph_nas.py b/autogl/module/nas/space/graph_nas.py index 42ce952..57efedb 100644 --- a/autogl/module/nas/space/graph_nas.py +++ b/autogl/module/nas/space/graph_nas.py @@ -30,9 +30,14 @@ GRAPHNAS_DEFAULT_GNN_OPS = [ GRAPHNAS_DEFAULT_ACT_OPS = [ # "sigmoid", "tanh", "relu", "linear", # "softplus", "leaky_relu", "relu6", "elu" - "sigmoid", "tanh", "relu", "linear", "elu" + "sigmoid", + "tanh", + "relu", + "linear", + "elu", ] + class LambdaModule(nn.Module): def __init__(self, lambd): super().__init__() @@ -40,27 +45,31 @@ class LambdaModule(nn.Module): def forward(self, x): return self.lambd(x) - + def __repr__(self): - return '{}({})'.format(self.__class__.__name__,self.lambd) + return "{}({})".format(self.__class__.__name__, self.lambd) + class StrModule(nn.Module): def __init__(self, lambd): super().__init__() self.str = lambd - def forward(self, *args,**kwargs): - return self.str + def forward(self, *args, **kwargs): + return self.str def __repr__(self): - return '{}({})'.format(self.__class__.__name__,self.str) + return "{}({})".format(self.__class__.__name__, self.str) + def act_map_nn(act): return LambdaModule(act_map(act)) + def map_nn(l): return [StrModule(x) for x in l] + @register_nas_space("graphnas") class GraphNasNodeClassificationSpace(BaseSpace): def __init__( @@ -71,7 +80,7 @@ class GraphNasNodeClassificationSpace(BaseSpace): input_dim: _typ.Optional[int] = None, output_dim: _typ.Optional[int] = None, gnn_ops: _typ.Sequence[_typ.Union[str, _typ.Any]] = GRAPHNAS_DEFAULT_GNN_OPS, - act_ops: _typ.Sequence[_typ.Union[str, _typ.Any]] = GRAPHNAS_DEFAULT_ACT_OPS + act_ops: _typ.Sequence[_typ.Union[str, _typ.Any]] = GRAPHNAS_DEFAULT_ACT_OPS, ): super().__init__() self.layer_number = layer_number @@ -81,7 +90,7 @@ class GraphNasNodeClassificationSpace(BaseSpace): self.gnn_ops = gnn_ops self.act_ops = act_ops self.dropout = dropout - + def instantiate( self, hidden_dim: _typ.Optional[int] = None, @@ -90,7 +99,7 @@ class GraphNasNodeClassificationSpace(BaseSpace): input_dim: _typ.Optional[int] = None, output_dim: _typ.Optional[int] = None, gnn_ops: _typ.Sequence[_typ.Union[str, _typ.Any]] = None, - act_ops: _typ.Sequence[_typ.Union[str, _typ.Any]] = None + act_ops: _typ.Sequence[_typ.Union[str, _typ.Any]] = None, ): super().instantiate() self.dropout = dropout or self.dropout @@ -103,45 +112,80 @@ class GraphNasNodeClassificationSpace(BaseSpace): self.preproc0 = nn.Linear(self.input_dim, self.hidden_dim) self.preproc1 = nn.Linear(self.input_dim, self.hidden_dim) node_labels = [mutables.InputChoice.NO_KEY, mutables.InputChoice.NO_KEY] - for layer in range(2,self.layer_number+2): + for layer in range(2, self.layer_number + 2): node_labels.append(f"op_{layer}") - setattr(self,f"in_{layer}",self.setInputChoice(layer,choose_from=node_labels[:-1], n_chosen=1, return_mask=False,key=f"in_{layer}")) - setattr(self,f"op_{layer}",self.setLayerChoice(layer,[gnn_map(op,self.hidden_dim,self.hidden_dim)for op in self.gnn_ops],key=f"op_{layer}")) - setattr(self,"act",self.setLayerChoice(2*layer,[act_map_nn(a)for a in self.act_ops],key="act")) - setattr(self,"concat",self.setLayerChoice(2*layer+1,map_nn(["add", "product", "concat"]) ,key="concat")) + setattr( + self, + f"in_{layer}", + self.setInputChoice( + layer, + choose_from=node_labels[:-1], + n_chosen=1, + return_mask=False, + key=f"in_{layer}", + ), + ) + setattr( + self, + f"op_{layer}", + self.setLayerChoice( + layer, + [ + gnn_map(op, self.hidden_dim, self.hidden_dim) + for op in self.gnn_ops + ], + key=f"op_{layer}", + ), + ) + setattr( + self, + "act", + self.setLayerChoice( + 2 * layer, [act_map_nn(a) for a in self.act_ops], key="act" + ), + ) + setattr( + self, + "concat", + self.setLayerChoice( + 2 * layer + 1, map_nn(["add", "product", "concat"]), key="concat" + ), + ) self._initialized = True - self.classifier1 = nn.Linear(self.hidden_dim*self.layer_number, self.output_dim) + self.classifier1 = nn.Linear( + self.hidden_dim * self.layer_number, self.output_dim + ) self.classifier2 = nn.Linear(self.hidden_dim, self.output_dim) def forward(self, data): - x, edges = data.x, data.edge_index # x [2708,1433] ,[2, 10556] - x = F.dropout(x, p=self.dropout, training = self.training) + x, edges = data.x, data.edge_index # x [2708,1433] ,[2, 10556] + x = F.dropout(x, p=self.dropout, training=self.training) pprev_, prev_ = self.preproc0(x), self.preproc1(x) - prev_nodes_out = [pprev_,prev_] - for layer in range(2,self.layer_number+2): + prev_nodes_out = [pprev_, prev_] + for layer in range(2, self.layer_number + 2): node_in = getattr(self, f"in_{layer}")(prev_nodes_out) - node_out= getattr(self, f"op_{layer}")(node_in,edges) + node_out = getattr(self, f"op_{layer}")(node_in, edges) prev_nodes_out.append(node_out) - act=getattr(self, "act") - con=getattr(self, "concat")() - states=prev_nodes_out + act = getattr(self, "act") + con = getattr(self, "concat")() + states = prev_nodes_out if con == "concat": - x=torch.cat(states[2:], dim=1) + x = torch.cat(states[2:], dim=1) else: tmp = states[2] - for i in range(2,len(states)): + for i in range(2, len(states)): if con == "add": tmp = torch.add(tmp, states[i]) elif con == "product": tmp = torch.mul(tmp, states[i]) - x=tmp + x = tmp x = act(x) - if con=='concat': - x=self.classifier1(x) + if con == "concat": + x = self.classifier1(x) else: - x=self.classifier2(x) + x = self.classifier2(x) return F.log_softmax(x, dim=1) def parse_model(self, selection, device) -> BaseModel: - #return AutoGCN(self.input_dim, self.output_dim, device) - return self.wrap(device).fix(selection) \ No newline at end of file + # return AutoGCN(self.input_dim, self.output_dim, device) + return self.wrap(device).fix(selection) diff --git a/autogl/module/nas/space/graph_nas_macro.py b/autogl/module/nas/space/graph_nas_macro.py index ddca04a..4929399 100644 --- a/autogl/module/nas/space/graph_nas_macro.py +++ b/autogl/module/nas/space/graph_nas_macro.py @@ -10,7 +10,12 @@ from .operation import act_map from torch.nn import Parameter from torch_geometric.nn.inits import glorot, zeros -from torch_geometric.utils import remove_self_loops, add_self_loops, add_remaining_self_loops, softmax +from torch_geometric.utils import ( + remove_self_loops, + add_self_loops, + add_remaining_self_loops, + softmax, +) from torch_scatter import scatter_add import torch_scatter @@ -18,14 +23,22 @@ import inspect import sys special_args = [ - 'edge_index', 'edge_index_i', 'edge_index_j', 'size', 'size_i', 'size_j' + "edge_index", + "edge_index_i", + "edge_index_j", + "size", + "size_i", + "size_j", ] -__size_error_msg__ = ('All tensors which should get mapped to the same source ' - 'or target nodes must be of same size in dimension 0.') +__size_error_msg__ = ( + "All tensors which should get mapped to the same source " + "or target nodes must be of same size in dimension 0." +) is_python2 = sys.version_info[0] < 3 getargspec = inspect.getargspec if is_python2 else inspect.getfullargspec + def scatter_(name, src, index, dim_size=None): r"""Aggregates all values from the :attr:`src` tensor at the indices specified in the :attr:`index` tensor along the first dimension. @@ -45,35 +58,37 @@ def scatter_(name, src, index, dim_size=None): :rtype: :class:`Tensor` """ - assert name in ['add', 'mean', 'max'] + assert name in ["add", "mean", "max"] - op = getattr(torch_scatter, 'scatter_{}'.format(name)) - fill_value = -1e9 if name == 'max' else 0 + op = getattr(torch_scatter, "scatter_{}".format(name)) + fill_value = -1e9 if name == "max" else 0 out = op(src, index, 0, None, dim_size) if isinstance(out, tuple): out = out[0] - if name == 'max': + if name == "max": out[out == fill_value] = 0 return out -class MessagePassing(torch.nn.Module): - def __init__(self, aggr='add', flow='source_to_target'): +class MessagePassing(torch.nn.Module): + def __init__(self, aggr="add", flow="source_to_target"): super(MessagePassing, self).__init__() self.aggr = aggr - assert self.aggr in ['add', 'mean', 'max'] + assert self.aggr in ["add", "mean", "max"] self.flow = flow - assert self.flow in ['source_to_target', 'target_to_source'] + assert self.flow in ["source_to_target", "target_to_source"] self.__message_args__ = getargspec(self.message)[0][1:] - self.__special_args__ = [(i, arg) - for i, arg in enumerate(self.__message_args__) - if arg in special_args] + self.__special_args__ = [ + (i, arg) + for i, arg in enumerate(self.__message_args__) + if arg in special_args + ] self.__message_args__ = [ arg for arg in self.__message_args__ if arg not in special_args ] @@ -96,7 +111,7 @@ class MessagePassing(torch.nn.Module): size = [None, None] if size is None else list(size) assert len(size) == 2 - i, j = (0, 1) if self.flow == 'target_to_source' else (1, 0) + i, j = (0, 1) if self.flow == "target_to_source" else (1, 0) ij = {"_i": i, "_j": j} message_args = [] @@ -129,8 +144,8 @@ class MessagePassing(torch.nn.Module): size[0] = size[1] if size[0] is None else size[0] size[1] = size[0] if size[1] is None else size[1] - kwargs['edge_index'] = edge_index - kwargs['size'] = size + kwargs["edge_index"] = edge_index + kwargs["size"] = size for (idx, arg) in self.__special_args__: if arg[-2:] in ij.keys(): @@ -168,21 +183,23 @@ class MessagePassing(torch.nn.Module): return aggr_out -class GeoLayer(MessagePassing): - def __init__(self, - in_channels, - out_channels, - heads=1, - concat=True, - negative_slope=0.2, - dropout=0, - bias=True, - att_type="gat", - agg_type="sum", - pool_dim=0): +class GeoLayer(MessagePassing): + def __init__( + self, + in_channels, + out_channels, + heads=1, + concat=True, + negative_slope=0.2, + dropout=0, + bias=True, + att_type="gat", + agg_type="sum", + pool_dim=0, + ): if agg_type in ["sum", "mlp"]: - super(GeoLayer, self).__init__('add') + super(GeoLayer, self).__init__("add") elif agg_type in ["mean", "max"]: super(GeoLayer, self).__init__(agg_type) self.in_channels = in_channels @@ -197,8 +214,7 @@ class GeoLayer(MessagePassing): # GCN weight self.gcn_weight = None - self.weight = Parameter( - torch.Tensor(in_channels, heads * out_channels)) + self.weight = Parameter(torch.Tensor(in_channels, heads * out_channels)) self.att = Parameter(torch.Tensor(1, heads, 2 * out_channels)) if bias and concat: @@ -206,7 +222,7 @@ class GeoLayer(MessagePassing): elif bias and not concat: self.bias = Parameter(torch.Tensor(out_channels)) else: - self.register_parameter('bias', None) + self.register_parameter("bias", None) if self.att_type in ["generalized_linear"]: self.general_att_layer = torch.nn.Linear(out_channels, 1, bias=False) @@ -226,18 +242,19 @@ class GeoLayer(MessagePassing): @staticmethod def norm(edge_index, num_nodes, edge_weight, improved=False, dtype=None): if edge_weight is None: - edge_weight = torch.ones((edge_index.size(1), ), - dtype=dtype, - device=edge_index.device) + edge_weight = torch.ones( + (edge_index.size(1),), dtype=dtype, device=edge_index.device + ) fill_value = 1 if not improved else 2 edge_index, edge_weight = add_remaining_self_loops( - edge_index, edge_weight, fill_value, num_nodes) + edge_index, edge_weight, fill_value, num_nodes + ) row, col = edge_index deg = scatter_add(edge_weight, row, dim=0, dim_size=num_nodes) deg_inv_sqrt = deg.pow(-0.5) - deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0 + deg_inv_sqrt[deg_inv_sqrt == float("inf")] = 0 return edge_index, deg_inv_sqrt[row] * edge_weight * deg_inv_sqrt[col] @@ -269,14 +286,16 @@ class GeoLayer(MessagePassing): x_j = F.dropout(x_j, p=self.dropout, training=True) neighbor = x_j elif self.att_type == "gcn": - if self.gcn_weight is None or self.gcn_weight.size(0) != x_j.size(0): # 对于不同的图gcn_weight需要重新计算 + if self.gcn_weight is None or self.gcn_weight.size(0) != x_j.size( + 0 + ): # 对于不同的图gcn_weight需要重新计算 _, norm = self.norm(edge_index, num_nodes, None) self.gcn_weight = norm neighbor = self.gcn_weight.view(-1, 1, 1) * x_j else: # Compute attention coefficients. alpha = self.apply_attention(edge_index, num_nodes, x_i, x_j) - alpha = softmax(alpha, edge_index[0], num_nodes = num_nodes) + alpha = softmax(alpha, edge_index[0], num_nodes=num_nodes) # Sample attention coefficients stochastically. if self.training and self.dropout > 0: alpha = F.dropout(alpha, p=self.dropout, training=True) @@ -293,28 +312,30 @@ class GeoLayer(MessagePassing): alpha = F.leaky_relu(alpha, self.negative_slope) elif self.att_type == "gat_sym": - wl = self.att[:, :, :self.out_channels] # weight left - wr = self.att[:, :, self.out_channels:] # weight right + wl = self.att[:, :, : self.out_channels] # weight left + wr = self.att[:, :, self.out_channels :] # weight right alpha = (x_i * wl).sum(dim=-1) + (x_j * wr).sum(dim=-1) alpha_2 = (x_j * wl).sum(dim=-1) + (x_i * wr).sum(dim=-1) - alpha = F.leaky_relu(alpha, self.negative_slope) + F.leaky_relu(alpha_2, self.negative_slope) + alpha = F.leaky_relu(alpha, self.negative_slope) + F.leaky_relu( + alpha_2, self.negative_slope + ) elif self.att_type == "linear": - wl = self.att[:, :, :self.out_channels] # weight left - wr = self.att[:, :, self.out_channels:] # weight right + wl = self.att[:, :, : self.out_channels] # weight left + wr = self.att[:, :, self.out_channels :] # weight right al = x_j * wl ar = x_j * wr alpha = al.sum(dim=-1) + ar.sum(dim=-1) alpha = torch.tanh(alpha) elif self.att_type == "cos": - wl = self.att[:, :, :self.out_channels] # weight left - wr = self.att[:, :, self.out_channels:] # weight right + wl = self.att[:, :, : self.out_channels] # weight left + wr = self.att[:, :, self.out_channels :] # weight right alpha = x_i * wl * x_j * wr alpha = alpha.sum(dim=-1) elif self.att_type == "generalized_linear": - wl = self.att[:, :, :self.out_channels] # weight left - wr = self.att[:, :, self.out_channels:] # weight right + wl = self.att[:, :, : self.out_channels] # weight left + wr = self.att[:, :, self.out_channels :] # weight right al = x_i * wl ar = x_j * wr alpha = al + ar @@ -335,9 +356,9 @@ class GeoLayer(MessagePassing): return aggr_out def __repr__(self): - return '{}({}, {}, heads={})'.format(self.__class__.__name__, - self.in_channels, - self.out_channels, self.heads) + return "{}({}, {}, heads={})".format( + self.__class__.__name__, self.in_channels, self.out_channels, self.heads + ) def get_param_dict(self): params = {} @@ -374,6 +395,7 @@ class GeoLayer(MessagePassing): if agg_key in params and hasattr(self, "pool_layer"): self.pool_layer.load_state_dict(params[agg_key]) + @register_nas_space("graphnasmacro") class GraphNasMacroNodeClassificationSpace(BaseSpace): def __init__( @@ -384,7 +406,7 @@ class GraphNasMacroNodeClassificationSpace(BaseSpace): input_dim: _typ.Optional[int] = None, output_dim: _typ.Optional[int] = None, ops: _typ.Tuple = None, - search_act_con=False + search_act_con=False, ): super().__init__() self.layer_number = layer_number @@ -393,7 +415,7 @@ class GraphNasMacroNodeClassificationSpace(BaseSpace): self.output_dim = output_dim self.ops = ops self.dropout = dropout - self.search_act_con=search_act_con + self.search_act_con = search_act_con def instantiate( self, @@ -402,7 +424,7 @@ class GraphNasMacroNodeClassificationSpace(BaseSpace): input_dim: _typ.Optional[int] = None, output_dim: _typ.Optional[int] = None, ops: _typ.Tuple = None, - dropout = None + dropout=None, ): super().instantiate() self.hidden_dim = hidden_dim or self.hidden_dim @@ -421,32 +443,145 @@ class GraphNasMacroNodeClassificationSpace(BaseSpace): # build hidden layer for i in range(layer_nums): # extract layer information - setattr(self,f"attention_{i}",self.setLayerChoice(i * state_num + 0, map_nn(["gat", "gcn", "cos", "const", "gat_sym", 'linear', 'generalized_linear']), key = f"attention_{i}")) - setattr(self,f"aggregator_{i}",self.setLayerChoice(i * state_num + 1, map_nn(["sum", "mean", "max", "mlp", ]), key = f"aggregator_{i}")) - setattr(self,f"act_{i}",self.setLayerChoice(i * state_num + 0, map_nn(["sigmoid", "tanh", "relu", "linear", - "softplus", "leaky_relu", "relu6", "elu"]), key=f"act_{i}")) - setattr(self,f"head_{i}",self.setLayerChoice(i * state_num + 0, map_nn([1, 2, 4, 6, 8, 16]), key= f"head_{i}")) + setattr( + self, + f"attention_{i}", + self.setLayerChoice( + i * state_num + 0, + map_nn( + [ + "gat", + "gcn", + "cos", + "const", + "gat_sym", + "linear", + "generalized_linear", + ] + ), + key=f"attention_{i}", + ), + ) + setattr( + self, + f"aggregator_{i}", + self.setLayerChoice( + i * state_num + 1, + map_nn( + [ + "sum", + "mean", + "max", + "mlp", + ] + ), + key=f"aggregator_{i}", + ), + ) + setattr( + self, + f"act_{i}", + self.setLayerChoice( + i * state_num + 0, + map_nn( + [ + "sigmoid", + "tanh", + "relu", + "linear", + "softplus", + "leaky_relu", + "relu6", + "elu", + ] + ), + key=f"act_{i}", + ), + ) + setattr( + self, + f"head_{i}", + self.setLayerChoice( + i * state_num + 0, map_nn([1, 2, 4, 6, 8, 16]), key=f"head_{i}" + ), + ) if i < layer_nums - 1: - setattr(self,f"out_channels_{i}",self.setLayerChoice(i * state_num + 0, map_nn([4, 8, 16, 32, 64, 128, 256]), key=f"out_channels_{i}")) + setattr( + self, + f"out_channels_{i}", + self.setLayerChoice( + i * state_num + 0, + map_nn([4, 8, 16, 32, 64, 128, 256]), + key=f"out_channels_{i}", + ), + ) def parse_model(self, selection, device) -> BaseModel: sel_list = [] for i in range(self.layer_number): - sel_list.append(["gat", "gcn", "cos", "const", "gat_sym", 'linear', 'generalized_linear'][selection[f"attention_{i}"]]) - sel_list.append(["sum", "mean", "max", "mlp", ][selection[f"aggregator_{i}"]]) - sel_list.append(["sigmoid", "tanh", "relu", "linear","softplus", "leaky_relu", "relu6", "elu"][selection[f"act_{i}"]]) + sel_list.append( + [ + "gat", + "gcn", + "cos", + "const", + "gat_sym", + "linear", + "generalized_linear", + ][selection[f"attention_{i}"]] + ) + sel_list.append( + [ + "sum", + "mean", + "max", + "mlp", + ][selection[f"aggregator_{i}"]] + ) + sel_list.append( + [ + "sigmoid", + "tanh", + "relu", + "linear", + "softplus", + "leaky_relu", + "relu6", + "elu", + ][selection[f"act_{i}"]] + ) sel_list.append([1, 2, 4, 6, 8, 16][selection[f"head_{i}"]]) if i < self.layer_number - 1: - sel_list.append([4, 8, 16, 32, 64, 128, 256][selection[f"out_channels_{i}"]]) + sel_list.append( + [4, 8, 16, 32, 64, 128, 256][selection[f"out_channels_{i}"]] + ) sel_list.append(self.output_dim) - #sel_list = ['const', 'sum', 'relu6', 2, 128, 'gat', 'sum', 'linear', 2, 7] - model = GraphNet(sel_list, self.input_dim, self.output_dim, self.dropout, multi_label=False, batch_normal=False, layers = self.layer_number).wrap(device) + # sel_list = ['const', 'sum', 'relu6', 2, 128, 'gat', 'sum', 'linear', 2, 7] + model = GraphNet( + sel_list, + self.input_dim, + self.output_dim, + self.dropout, + multi_label=False, + batch_normal=False, + layers=self.layer_number, + ).wrap(device) return model -class GraphNet(BaseSpace): - def __init__(self, actions, num_feat, num_label, drop_out=0.6, multi_label=False, batch_normal=True, state_num=5, - residual=False, layers = 2): +class GraphNet(BaseSpace): + def __init__( + self, + actions, + num_feat, + num_label, + drop_out=0.6, + multi_label=False, + batch_normal=True, + state_num=5, + residual=False, + layers=2, + ): self.residual = residual self.batch_normal = batch_normal self.layer_nums = layers @@ -456,11 +591,15 @@ class GraphNet(BaseSpace): self.input_dim = num_feat self.output_dim = num_label self.dropout = drop_out - + super().__init__() - self.build_model(actions, batch_normal, drop_out, num_feat, num_label, state_num) + self.build_model( + actions, batch_normal, drop_out, num_feat, num_label, state_num + ) - def build_model(self, actions, batch_normal, drop_out, num_feat, num_label, state_num): + def build_model( + self, actions, batch_normal, drop_out, num_feat, num_label, state_num + ): if self.residual: self.fcs = torch.nn.ModuleList() if self.batch_normal: @@ -468,9 +607,26 @@ class GraphNet(BaseSpace): self.layers = torch.nn.ModuleList() self.acts = [] self.gates = torch.nn.ModuleList() - self.build_hidden_layers(actions, batch_normal, drop_out, self.layer_nums, num_feat, num_label, state_num) - - def build_hidden_layers(self, actions, batch_normal, drop_out, layer_nums, num_feat, num_label, state_num=6): + self.build_hidden_layers( + actions, + batch_normal, + drop_out, + self.layer_nums, + num_feat, + num_label, + state_num, + ) + + def build_hidden_layers( + self, + actions, + batch_normal, + drop_out, + layer_nums, + num_feat, + num_label, + state_num=6, + ): # build hidden layer for i in range(layer_nums): @@ -492,17 +648,27 @@ class GraphNet(BaseSpace): if self.batch_normal: self.bns.append(torch.nn.BatchNorm1d(in_channels, momentum=0.5)) self.layers.append( - GeoLayer(in_channels, out_channels, head_num, concat, dropout=self.dropout, - att_type=attention_type, agg_type=aggregator_type, )) + GeoLayer( + in_channels, + out_channels, + head_num, + concat, + dropout=self.dropout, + att_type=attention_type, + agg_type=aggregator_type, + ) + ) self.acts.append(act_map(act)) if self.residual: if concat: - self.fcs.append(torch.nn.Linear(in_channels, out_channels * head_num)) + self.fcs.append( + torch.nn.Linear(in_channels, out_channels * head_num) + ) else: self.fcs.append(torch.nn.Linear(in_channels, out_channels)) def forward(self, data): - output, edge_index_all = data.x, data.edge_index # x [2708,1433] ,[2, 10556] + output, edge_index_all = data.x, data.edge_index # x [2708,1433] ,[2, 10556] if self.residual: for i, (act, layer, fc) in enumerate(zip(self.acts, self.layers, self.fcs)): output = F.dropout(output, p=self.dropout, training=self.training) diff --git a/autogl/module/nas/space/operation.py b/autogl/module/nas/space/operation.py index e642928..1690305 100644 --- a/autogl/module/nas/space/operation.py +++ b/autogl/module/nas/space/operation.py @@ -14,6 +14,7 @@ import torch from torch import nn import torch.nn.functional as F + class LinearConv(nn.Module): def __init__(self, in_channels, out_channels, bias=True): super(LinearConv, self).__init__() @@ -26,24 +27,28 @@ class LinearConv(nn.Module): return self.linear(x) def __repr__(self): - return '{}({}, {})'.format(self.__class__.__name__, self.in_channels, - self.out_channels) + return "{}({}, {})".format( + self.__class__.__name__, self.in_channels, self.out_channels + ) -class ZeroConv(nn.Module): +class ZeroConv(nn.Module): def forward(self, x, edge_index, edge_weight=None): out = torch.zeros_like(x) out.requires_grad = True return out def __repr__(self): - return 'ZeroConv()' + return "ZeroConv()" + class Identity(nn.Module): def forward(self, x, edge_index, edge_weight=None): return x + def __repr__(self): - return 'Identity()' + return "Identity()" + def act_map(act): if act == "linear": @@ -65,15 +70,16 @@ def act_map(act): else: raise Exception("wrong activate function") + def gnn_map(gnn_name, in_dim, out_dim, concat=False, bias=True) -> nn.Module: - ''' + """ :param gnn_name: :param in_dim: :param out_dim: :param concat: for gat, concat multi-head output or not :return: GNN model - ''' + """ if gnn_name == "gat_8": return GATConv(in_dim, out_dim, 8, concat=concat, bias=bias) elif gnn_name == "gat_6": @@ -100,12 +106,21 @@ def gnn_map(gnn_name, in_dim, out_dim, concat=False, bias=True) -> nn.Module: return LinearConv(in_dim, out_dim, bias=bias) elif gnn_name == "zero": return ZeroConv() - elif gnn_name == 'identity': + elif gnn_name == "identity": return Identity() elif hasattr(torch_geometric.nn, gnn_name): cls = getattr(torch_geometric.nn, gnn_name) assert isinstance(cls, type), "Only support modules, get %s" % (gnn_name) - kwargs = {'in_channels': in_dim, 'out_channels': out_dim, 'concat': concat, 'bias': bias} - kwargs = {key: kwargs[key] for key in cls.__init__.__code__.co_varnames if key in kwargs} + kwargs = { + "in_channels": in_dim, + "out_channels": out_dim, + "concat": concat, + "bias": bias, + } + kwargs = { + key: kwargs[key] + for key in cls.__init__.__code__.co_varnames + if key in kwargs + } return cls(**kwargs) raise KeyError("Cannot parse key %s" % (gnn_name)) diff --git a/autogl/module/nas/space/single_path.py b/autogl/module/nas/space/single_path.py index 3475b42..eadff77 100644 --- a/autogl/module/nas/space/single_path.py +++ b/autogl/module/nas/space/single_path.py @@ -12,6 +12,7 @@ from ....utils import get_logger from ...model import AutoGCN + @register_nas_space("singlepath") class SinglePathNodeClassificationSpace(BaseSpace): def __init__( @@ -21,7 +22,7 @@ class SinglePathNodeClassificationSpace(BaseSpace): dropout: _typ.Optional[float] = 0.2, input_dim: _typ.Optional[int] = None, output_dim: _typ.Optional[int] = None, - ops: _typ.Tuple = ['GCNConv', 'GATConv'], + ops: _typ.Tuple = ["GCNConv", "GATConv"], ): super().__init__() self.layer_number = layer_number @@ -38,7 +39,7 @@ class SinglePathNodeClassificationSpace(BaseSpace): input_dim: _typ.Optional[int] = None, output_dim: _typ.Optional[int] = None, ops: _typ.Tuple = None, - dropout = None + dropout=None, ): super().instantiate() self.hidden_dim = hidden_dim or self.hidden_dim @@ -59,8 +60,15 @@ class SinglePathNodeClassificationSpace(BaseSpace): self.output_dim if layer == self.layer_number - 1 else self.hidden_dim, - ) if isinstance(op, type) else gnn_map(op, self.input_dim if layer == 0 else self.hidden_dim, - self.output_dim if layer == self.layer_number - 1 else self.hidden_dim) + ) + if isinstance(op, type) + else gnn_map( + op, + self.input_dim if layer == 0 else self.hidden_dim, + self.output_dim + if layer == self.layer_number - 1 + else self.hidden_dim, + ) for op in self.ops ], ), @@ -73,9 +81,9 @@ class SinglePathNodeClassificationSpace(BaseSpace): x = getattr(self, f"op_{layer}")(x, edges) if layer != self.layer_number - 1: x = F.leaky_relu(x) - x = F.dropout(x, p=self.dropout, training = self.training) + x = F.dropout(x, p=self.dropout, training=self.training) return F.log_softmax(x, dim=1) def parse_model(self, selection, device) -> BaseModel: - #return AutoGCN(self.input_dim, self.output_dim, device) + # return AutoGCN(self.input_dim, self.output_dim, device) return self.wrap(device).fix(selection) diff --git a/autogl/module/nas/utils.py b/autogl/module/nas/utils.py index 2504cfc..9ca8a05 100644 --- a/autogl/module/nas/utils.py +++ b/autogl/module/nas/utils.py @@ -123,22 +123,26 @@ class AverageMeter: fmtstr = "{name}: {avg" + self.fmt + "}" return fmtstr.format(**self.__dict__) + def get_module_order(root_module): key2order = {} + def apply(m): for name, child in m.named_children(): if isinstance(child, Mutable): - key2order[child.key] = child.order + key2order[child.key] = child.order else: apply(child) apply(root_module) return key2order + def sort_replaced_module(k2o, modules): - modules = sorted(modules, key = lambda x:k2o[x[0]]) + modules = sorted(modules, key=lambda x: k2o[x[0]]) return modules + def _replace_module_with_type(root_module, init_fn, type_name, modules): if modules is None: modules = [] diff --git a/autogl/module/train/evaluation.py b/autogl/module/train/evaluation.py index d2a390f..595df12 100644 --- a/autogl/module/train/evaluation.py +++ b/autogl/module/train/evaluation.py @@ -28,6 +28,7 @@ class Evaluation: class EvaluatorUtility: """ Auxiliary utilities for evaluation """ + class PredictionBatchCumulativeBuilder: """ Batch-cumulative builder for prediction @@ -37,22 +38,22 @@ class EvaluatorUtility: a batch-cumulative prediction collector `PredictionBatchCumulativeBuilder` is implemented for prediction in mini-batch manner. """ + def __init__(self): self.__indexes_in_integral_data: _typing.Optional[np.ndarray] = None self.__prediction: _typing.Optional[np.ndarray] = None def clear_batches( - self, *__args, **__kwargs - ) -> 'EvaluatorUtility.PredictionBatchCumulativeBuilder': + self, *__args, **__kwargs + ) -> "EvaluatorUtility.PredictionBatchCumulativeBuilder": self.__indexes_in_integral_data = None self.__prediction = None return self def add_batch( - self, indexes_in_integral_data: np.ndarray, - batch_prediction: np.ndarray - ) -> 'EvaluatorUtility.PredictionBatchCumulativeBuilder': - if not( + self, indexes_in_integral_data: np.ndarray, batch_prediction: np.ndarray + ) -> "EvaluatorUtility.PredictionBatchCumulativeBuilder": + if not ( isinstance(indexes_in_integral_data, np.ndarray) and isinstance(batch_prediction, np.ndarray) and len(indexes_in_integral_data.shape) == 1 @@ -62,33 +63,50 @@ class EvaluatorUtility: raise ValueError if self.__indexes_in_integral_data is None: - if indexes_in_integral_data.shape != np.unique(indexes_in_integral_data).shape: + if ( + indexes_in_integral_data.shape + != np.unique(indexes_in_integral_data).shape + ): raise ValueError( f"There exists duplicate index " f"in the argument indexes_in_integral_data {indexes_in_integral_data}" ) else: - self.__indexes_in_integral_data: np.ndarray = np.unique(indexes_in_integral_data) + self.__indexes_in_integral_data: np.ndarray = np.unique( + indexes_in_integral_data + ) else: __indexes_in_integral_data = np.concatenate( (self.__indexes_in_integral_data, indexes_in_integral_data) ) - if __indexes_in_integral_data.shape != np.unique(__indexes_in_integral_data).shape: + if ( + __indexes_in_integral_data.shape + != np.unique(__indexes_in_integral_data).shape + ): raise ValueError else: - self.__indexes_in_integral_data: np.ndarray = __indexes_in_integral_data + self.__indexes_in_integral_data: np.ndarray = ( + __indexes_in_integral_data + ) if self.__prediction is None: self.__prediction: np.ndarray = batch_prediction else: - self.__prediction: np.ndarray = np.concatenate((self.__prediction, batch_prediction)) + self.__prediction: np.ndarray = np.concatenate( + (self.__prediction, batch_prediction) + ) return self - def compose(self, __sorted: bool = True, **__kwargs) -> _typing.Tuple[np.ndarray, np.ndarray]: + def compose( + self, __sorted: bool = True, **__kwargs + ) -> _typing.Tuple[np.ndarray, np.ndarray]: if __sorted: sorted_index = np.argsort(self.__indexes_in_integral_data) - return self.__indexes_in_integral_data[sorted_index], self.__prediction[sorted_index] + return ( + self.__indexes_in_integral_data[sorted_index], + self.__prediction[sorted_index], + ) else: return self.__indexes_in_integral_data, self.__prediction @@ -260,4 +278,4 @@ class MicroF1(Evaluation): @staticmethod def evaluate(predict, label) -> float: - return f1_score(label, np.argmax(predict, axis=1), average='micro') + return f1_score(label, np.argmax(predict, axis=1), average="micro") diff --git a/autogl/module/train/graph_classification_full.py b/autogl/module/train/graph_classification_full.py index 7504da8..8b12a05 100644 --- a/autogl/module/train/graph_classification_full.py +++ b/autogl/module/train/graph_classification_full.py @@ -399,14 +399,17 @@ class GraphClassificationFullTrainer(BaseGraphClassificationTrainer): def __repr__(self) -> str: import yaml - return yaml.dump({ - "trainer_name": self.__class__.__name__, - "optimizer": self.optimizer, - "learning_rate": self.lr, - "max_epoch": self.max_epoch, - "early_stopping_round": self.early_stopping_round, - "model": repr(self.model) - }) + + return yaml.dump( + { + "trainer_name": self.__class__.__name__, + "optimizer": self.optimizer, + "learning_rate": self.lr, + "max_epoch": self.max_epoch, + "early_stopping_round": self.early_stopping_round, + "model": repr(self.model), + } + ) def evaluate(self, dataset, mask="val", feval=None): """ diff --git a/autogl/module/train/node_classification_full.py b/autogl/module/train/node_classification_full.py index 65c1f6f..e9ec45d 100644 --- a/autogl/module/train/node_classification_full.py +++ b/autogl/module/train/node_classification_full.py @@ -57,7 +57,7 @@ class NodeClassificationFullTrainer(BaseNodeClassificationTrainer): def __init__( self, - model: Union[BaseModel, str]=None, + model: Union[BaseModel, str] = None, num_features=None, num_classes=None, optimizer=None, @@ -375,6 +375,7 @@ class NodeClassificationFullTrainer(BaseNodeClassificationTrainer): def __repr__(self) -> str: import yaml + return yaml.dump( { "trainer_name": self.__class__.__name__, @@ -382,7 +383,7 @@ class NodeClassificationFullTrainer(BaseNodeClassificationTrainer): "learning_rate": self.lr, "max_epoch": self.max_epoch, "early_stopping_round": self.early_stopping_round, - "model": repr(self.model) + "model": repr(self.model), } ) diff --git a/autogl/module/train/node_classification_trainer/node_classification_sampled_trainer.py b/autogl/module/train/node_classification_trainer/node_classification_sampled_trainer.py index 415e9bd..110a6c4 100644 --- a/autogl/module/train/node_classification_trainer/node_classification_sampled_trainer.py +++ b/autogl/module/train/node_classification_trainer/node_classification_sampled_trainer.py @@ -13,7 +13,7 @@ from ..sampling.sampler.target_dependant_sampler import TargetDependantSampledDa from ..sampling.sampler.neighbor_sampler import NeighborSampler from ..sampling.sampler.graphsaint_sampler import * from ..sampling.sampler.layer_dependent_importance_sampler import ( - LayerDependentImportanceSampler + LayerDependentImportanceSampler, ) from ...model import BaseModel from ...model.base import ClassificationSupportedSequentialModel @@ -28,14 +28,18 @@ class _DeterministicNeighborSamplerStore: ] = [] @classmethod - def __is_target_node_indexes_equal(cls, a: torch.LongTensor, b: torch.LongTensor) -> bool: + def __is_target_node_indexes_equal( + cls, a: torch.LongTensor, b: torch.LongTensor + ) -> bool: if not a.dtype == b.dtype == torch.int64: return False if a.size() != b.size(): return False return torch.where(a != b)[0].size(0) == 0 - def __setitem__(self, target_nodes: torch.Tensor, neighbor_sampler: NeighborSampler): + def __setitem__( + self, target_nodes: torch.Tensor, neighbor_sampler: NeighborSampler + ): target_nodes: _typing.Any = target_nodes.cpu() if type(target_nodes) != torch.Tensor or target_nodes.dtype != torch.int64: raise TypeError @@ -43,18 +47,25 @@ class _DeterministicNeighborSamplerStore: raise TypeError for i in range(len(self.__neighbor_sampler_mapping)): if self.__is_target_node_indexes_equal( - target_nodes, self.__neighbor_sampler_mapping[i][0] + target_nodes, self.__neighbor_sampler_mapping[i][0] ): self.__neighbor_sampler_mapping[i] = (target_nodes, neighbor_sampler) return self.__neighbor_sampler_mapping.append((target_nodes, neighbor_sampler)) - def __getitem__(self, target_nodes: torch.Tensor) -> _typing.Optional[NeighborSampler]: + def __getitem__( + self, target_nodes: torch.Tensor + ) -> _typing.Optional[NeighborSampler]: target_nodes: _typing.Any = target_nodes.cpu() if type(target_nodes) != torch.Tensor or target_nodes.dtype != torch.int64: raise TypeError - for __current_target_nodes, __neighbor_sampler in self.__neighbor_sampler_mapping: - if self.__is_target_node_indexes_equal(target_nodes, __current_target_nodes): + for ( + __current_target_nodes, + __neighbor_sampler, + ) in self.__neighbor_sampler_mapping: + if self.__is_target_node_indexes_equal( + target_nodes, __current_target_nodes + ): return __neighbor_sampler return None @@ -89,24 +100,25 @@ class NodeClassificationGraphSAINTTrainer(BaseNodeClassificationTrainer): feval: ``str``. The evaluation method adopted in this function. """ + def __init__( - self, - model: _typing.Union[BaseModel, str], - num_features: int, - num_classes: int, - optimizer: _typing.Union[_typing.Type[torch.optim.Optimizer], str, None] = ..., - lr: float = 1e-4, - max_epoch: int = 100, - early_stopping_round: int = 100, - weight_decay: float = 1e-4, - device: _typing.Optional[torch.device] = None, - init: bool = True, - feval: _typing.Union[ - _typing.Sequence[str], _typing.Sequence[_typing.Type[Evaluation]] - ] = (MicroF1,), - loss: str = "nll_loss", - lr_scheduler_type: _typing.Optional[str] = None, - **kwargs, + self, + model: _typing.Union[BaseModel, str], + num_features: int, + num_classes: int, + optimizer: _typing.Union[_typing.Type[torch.optim.Optimizer], str, None] = ..., + lr: float = 1e-4, + max_epoch: int = 100, + early_stopping_round: int = 100, + weight_decay: float = 1e-4, + device: _typing.Optional[torch.device] = None, + init: bool = True, + feval: _typing.Union[ + _typing.Sequence[str], _typing.Sequence[_typing.Type[Evaluation]] + ] = (MicroF1,), + loss: str = "nll_loss", + lr_scheduler_type: _typing.Optional[str] = None, + **kwargs, ): if isinstance(optimizer, type) and issubclass(optimizer, torch.optim.Optimizer): self._optimizer_class: _typing.Type[torch.optim.Optimizer] = optimizer @@ -137,10 +149,12 @@ class NodeClassificationGraphSAINTTrainer(BaseNodeClassificationTrainer): self._weight_decay: float = weight_decay if weight_decay > 0 else 1e-4 self._early_stopping = EarlyStopping( patience=early_stopping_round if early_stopping_round > 0 else 1e2, - verbose=False + verbose=False, ) """ Assign an empty initial hyper parameter space """ - self._hyper_parameter_space: _typing.Sequence[_typing.Dict[str, _typing.Any]] = [] + self._hyper_parameter_space: _typing.Sequence[ + _typing.Dict[str, _typing.Any] + ] = [] self._valid_result: torch.Tensor = torch.zeros(0) self._valid_result_prob: torch.Tensor = torch.zeros(0) @@ -200,13 +214,13 @@ class NodeClassificationGraphSAINTTrainer(BaseNodeClassificationTrainer): # self.__training_sampler_num_workers: int = kwargs.get( # "training_sampler_num_workers", _cpu_count() # ) - + # if not 0 <= self.__training_sampler_num_workers <= _cpu_count(): # self.__training_sampler_num_workers: int = _cpu_count() # force to be 0 to be compactible with current pyg solution. self.__training_sampler_num_workers: int = 0 - + super(NodeClassificationGraphSAINTTrainer, self).__init__( model, num_features, num_classes, device, init, feval, loss ) @@ -238,7 +252,7 @@ class NodeClassificationGraphSAINTTrainer(BaseNodeClassificationTrainer): optimizer: torch.optim.Optimizer = self._optimizer_class( self.model.model.parameters(), lr=self._learning_rate, - weight_decay=self._weight_decay + weight_decay=self._weight_decay, ) if type(self._lr_scheduler_type) == str: if self._lr_scheduler_type.lower() == "step" + "lr": @@ -269,42 +283,51 @@ class NodeClassificationGraphSAINTTrainer(BaseNodeClassificationTrainer): ) setattr( - integral_data, "edge_weight", - self.__compute_normalized_edge_weight(getattr(integral_data, "edge_index")) + integral_data, + "edge_weight", + self.__compute_normalized_edge_weight(getattr(integral_data, "edge_index")), ) " Generate Sampler " if self.__sampler_type.lower() == "edge": _sampler: torch_geometric.data.GraphSAINTEdgeSampler = ( GraphSAINTSamplerFactory.create_edge_sampler( - integral_data, self.__num_graphs_per_epoch, self.__sampled_budget, + integral_data, + self.__num_graphs_per_epoch, + self.__sampled_budget, self.__sample_coverage_factor, - num_workers=self.__training_sampler_num_workers + num_workers=self.__training_sampler_num_workers, ) ) elif self.__sampler_type.lower() == "node": _sampler: torch_geometric.data.GraphSAINTNodeSampler = ( GraphSAINTSamplerFactory.create_node_sampler( - integral_data, self.__num_graphs_per_epoch, self.__sampled_budget, + integral_data, + self.__num_graphs_per_epoch, + self.__sampled_budget, self.__sample_coverage_factor, - num_workers=self.__training_sampler_num_workers + num_workers=self.__training_sampler_num_workers, ) ) elif self.__sampler_type.lower() == "rw": _sampler: torch_geometric.data.GraphSAINTRandomWalkSampler = ( GraphSAINTSamplerFactory.create_random_walk_sampler( - integral_data, self.__num_graphs_per_epoch, - self.__sampled_budget, self.__walk_length, + integral_data, + self.__num_graphs_per_epoch, + self.__sampled_budget, + self.__walk_length, self.__sample_coverage_factor, - num_workers=self.__training_sampler_num_workers + num_workers=self.__training_sampler_num_workers, ) ) else: _sampler: torch_geometric.data.GraphSAINTRandomWalkSampler = ( GraphSAINTSamplerFactory.create_random_walk_sampler( - integral_data, self.__num_graphs_per_epoch, - self.__sampled_budget, self.__walk_length, + integral_data, + self.__num_graphs_per_epoch, + self.__sampled_budget, + self.__walk_length, self.__sample_coverage_factor, - num_workers=self.__training_sampler_num_workers + num_workers=self.__training_sampler_num_workers, ) ) @@ -315,31 +338,35 @@ class NodeClassificationGraphSAINTTrainer(BaseNodeClassificationTrainer): for sampled_data in _sampler: sampled_data = sampled_data.to(self.device) setattr( - sampled_data, "edge_weight", - getattr(sampled_data, "edge_norm") * getattr(sampled_data, "edge_weight") + sampled_data, + "edge_weight", + getattr(sampled_data, "edge_norm") + * getattr(sampled_data, "edge_weight"), ) optimizer.zero_grad() if isinstance(self.model.model, ClassificationSupportedSequentialModel): - prediction: torch.Tensor = self.model.model.cls_forward(sampled_data) + prediction: torch.Tensor = self.model.model.cls_forward( + sampled_data + ) else: prediction: torch.Tensor = self.model.model(sampled_data) if not hasattr(torch.nn.functional, self.loss): - raise TypeError( - f"PyTorch does not support loss type {self.loss}" - ) + raise TypeError(f"PyTorch does not support loss type {self.loss}") loss_function = getattr(torch.nn.functional, self.loss) loss_value: torch.Tensor = loss_function( - prediction, getattr(sampled_data, "y"), reduction='none' + prediction, getattr(sampled_data, "y"), reduction="none" ) - loss_value = (loss_value * getattr(sampled_data, "node_norm"))[sampled_data.train_mask].sum() + loss_value = (loss_value * getattr(sampled_data, "node_norm"))[ + sampled_data.train_mask + ].sum() loss_value.backward() optimizer.step() lr_scheduler.step() if ( - hasattr(integral_data, "val_mask") and - getattr(integral_data, "val_mask") is not None and - type(getattr(integral_data, "val_mask")) == torch.Tensor + hasattr(integral_data, "val_mask") + and getattr(integral_data, "val_mask") is not None + and type(getattr(integral_data, "val_mask")) == torch.Tensor ): validation_results: _typing.Sequence[float] = self.evaluate( (integral_data,), "val", [self.feval[0]] @@ -353,17 +380,16 @@ class NodeClassificationGraphSAINTTrainer(BaseNodeClassificationTrainer): LOGGER.debug("Early stopping at %d", current_epoch) break if ( - hasattr(integral_data, "val_mask") and - getattr(integral_data, "val_mask") is not None and - type(getattr(integral_data, "val_mask")) == torch.Tensor + hasattr(integral_data, "val_mask") + and getattr(integral_data, "val_mask") is not None + and type(getattr(integral_data, "val_mask")) == torch.Tensor ): self._early_stopping.load_checkpoint(self.model.model) def __predict_only( - self, integral_data, - mask_or_target_nodes_indexes: _typing.Union[ - torch.BoolTensor, torch.LongTensor - ] + self, + integral_data, + mask_or_target_nodes_indexes: _typing.Union[torch.BoolTensor, torch.LongTensor], ) -> torch.Tensor: """ The function of predicting on the given data. @@ -372,11 +398,13 @@ class NodeClassificationGraphSAINTTrainer(BaseNodeClassificationTrainer): :return: the result of prediction on the given dataset """ import copy + integral_data = copy.copy(integral_data) self.model.model.eval() setattr( - integral_data, "edge_weight", - self.__compute_normalized_edge_weight(getattr(integral_data, "edge_index")) + integral_data, + "edge_weight", + self.__compute_normalized_edge_weight(getattr(integral_data, "edge_index")), ) integral_data = integral_data.to(self.device) with torch.no_grad(): @@ -387,8 +415,7 @@ class NodeClassificationGraphSAINTTrainer(BaseNodeClassificationTrainer): return prediction[mask_or_target_nodes_indexes] def predict_proba( - self, dataset, mask: _typing.Optional[str] = None, - in_log_format: bool = False + self, dataset, mask: _typing.Optional[str] = None, in_log_format: bool = False ): """ The function of predicting the probability on the given dataset. @@ -443,12 +470,12 @@ class NodeClassificationGraphSAINTTrainer(BaseNodeClassificationTrainer): return self.predict_proba(dataset, mask, in_log_format=True).max(1)[1] def evaluate( - self, - dataset, - mask: _typing.Optional[str] = None, - feval: _typing.Union[ - None, _typing.Sequence[str], _typing.Sequence[_typing.Type[Evaluation]] - ] = None, + self, + dataset, + mask: _typing.Optional[str] = None, + feval: _typing.Union[ + None, _typing.Sequence[str], _typing.Sequence[_typing.Type[Evaluation]] + ] = None, ) -> _typing.Sequence[float]: """ The function of training on the given dataset and keeping valid result. @@ -492,13 +519,15 @@ class NodeClassificationGraphSAINTTrainer(BaseNodeClassificationTrainer): f.evaluate( prediction_probability.cpu().numpy(), y_ground_truth.cpu().numpy(), - ) for f in _feval + ) + for f in _feval ] @classmethod def __compute_normalized_edge_weight( - cls, edge_index: torch.LongTensor, - original_edge_weight: _typing.Optional[torch.Tensor] = ... + cls, + edge_index: torch.LongTensor, + original_edge_weight: _typing.Optional[torch.Tensor] = ..., ) -> torch.Tensor: if type(edge_index) != torch.Tensor: raise TypeError @@ -511,10 +540,8 @@ class NodeClassificationGraphSAINTTrainer(BaseNodeClassificationTrainer): elif original_edge_weight.size() != (edge_index.size(1),): original_edge_weight = original_edge_weight.resize(edge_index.size(1)) - __out_degree: torch.Tensor = \ - torch_geometric.utils.degree(edge_index[0]) - __in_degree: torch.Tensor = \ - torch_geometric.utils.degree(edge_index[1]) + __out_degree: torch.Tensor = torch_geometric.utils.degree(edge_index[0]) + __in_degree: torch.Tensor = torch_geometric.utils.degree(edge_index[1]) temp_tensor: torch.Tensor = torch.stack( [__out_degree[edge_index[0]], __in_degree[edge_index[1]]] ) @@ -539,6 +566,7 @@ class NodeClassificationGraphSAINTTrainer(BaseNodeClassificationTrainer): None """ import gc + gc.collect() data = dataset[0].to(torch.device("cpu")) self.__train_only(data) @@ -555,10 +583,10 @@ class NodeClassificationGraphSAINTTrainer(BaseNodeClassificationTrainer): return self._valid_result_prob def get_valid_score( - self, return_major: bool = True + self, return_major: bool = True ) -> _typing.Union[ _typing.Tuple[float, bool], - _typing.Tuple[_typing.Sequence[float], _typing.Sequence[bool]] + _typing.Tuple[_typing.Sequence[float], _typing.Sequence[bool]], ]: """ The function of getting the valid score. @@ -584,7 +612,7 @@ class NodeClassificationGraphSAINTTrainer(BaseNodeClassificationTrainer): @hyper_parameter_space.setter def hyper_parameter_space( - self, hp_space: _typing.Sequence[_typing.Dict[str, _typing.Any]] + self, hp_space: _typing.Sequence[_typing.Dict[str, _typing.Any]] ) -> None: if not isinstance(hp_space, _typing.Sequence): raise TypeError @@ -592,6 +620,7 @@ class NodeClassificationGraphSAINTTrainer(BaseNodeClassificationTrainer): def __repr__(self) -> str: import yaml + __repr: dict = { "trainer_name": self.__class__.__name__, "learning_rate": self._learning_rate, @@ -599,16 +628,16 @@ class NodeClassificationGraphSAINTTrainer(BaseNodeClassificationTrainer): "max_epoch": self._max_epoch, "early_stopping_round": self._early_stopping.patience, "sampler_type": self.__sampler_type, - "sampled_budget": self.__sampled_budget + "sampled_budget": self.__sampled_budget, } if self.__sampler_type == "rw": __repr.update({"walk_length": self.__walk_length}) return yaml.dump(__repr) def duplicate_from_hyper_parameter( - self, - hp: _typing.Dict[str, _typing.Any], - model: _typing.Optional[BaseModel] = None, + self, + hp: _typing.Dict[str, _typing.Any], + model: _typing.Optional[BaseModel] = None, ) -> "NodeClassificationGraphSAINTTrainer": """ The function of duplicating a new instance from the given hyper-parameter. @@ -646,12 +675,14 @@ class NodeClassificationGraphSAINTTrainer(BaseNodeClassificationTrainer): feval=self.feval, loss=self.loss, lr_scheduler_type=self._lr_scheduler_type, - **hp + **hp, ) @register_trainer("NodeClassificationLayerDependentImportanceSamplingTrainer") -class NodeClassificationLayerDependentImportanceSamplingTrainer(BaseNodeClassificationTrainer): +class NodeClassificationLayerDependentImportanceSamplingTrainer( + BaseNodeClassificationTrainer +): """ The node classification trainer utilizing Layer dependent importance sampling technique. @@ -680,24 +711,25 @@ class NodeClassificationLayerDependentImportanceSamplingTrainer(BaseNodeClassifi feval: ``str``. The evaluation method adopted in this function. """ + def __init__( - self, - model: _typing.Union[BaseModel, str], - num_features: int, - num_classes: int, - optimizer: _typing.Union[_typing.Type[torch.optim.Optimizer], str, None] = ..., - lr: float = 1e-4, - max_epoch: int = 100, - early_stopping_round: int = 100, - weight_decay: float = 1e-4, - device: _typing.Optional[torch.device] = None, - init: bool = True, - feval: _typing.Union[ - _typing.Sequence[str], _typing.Sequence[_typing.Type[Evaluation]] - ] = (MicroF1,), - loss: str = "nll_loss", - lr_scheduler_type: _typing.Optional[str] = None, - **kwargs, + self, + model: _typing.Union[BaseModel, str], + num_features: int, + num_classes: int, + optimizer: _typing.Union[_typing.Type[torch.optim.Optimizer], str, None] = ..., + lr: float = 1e-4, + max_epoch: int = 100, + early_stopping_round: int = 100, + weight_decay: float = 1e-4, + device: _typing.Optional[torch.device] = None, + init: bool = True, + feval: _typing.Union[ + _typing.Sequence[str], _typing.Sequence[_typing.Type[Evaluation]] + ] = (MicroF1,), + loss: str = "nll_loss", + lr_scheduler_type: _typing.Optional[str] = None, + **kwargs, ) -> None: if isinstance(optimizer, type) and issubclass(optimizer, torch.optim.Optimizer): self._optimizer_class: _typing.Type[torch.optim.Optimizer] = optimizer @@ -728,17 +760,21 @@ class NodeClassificationLayerDependentImportanceSamplingTrainer(BaseNodeClassifi self._weight_decay: float = weight_decay if weight_decay > 0 else 1e-4 self._early_stopping = EarlyStopping( patience=early_stopping_round if early_stopping_round > 0 else 1e2, - verbose=False + verbose=False, ) """ Assign an empty initial hyper parameter space """ - self._hyper_parameter_space: _typing.Sequence[_typing.Dict[str, _typing.Any]] = [] + self._hyper_parameter_space: _typing.Sequence[ + _typing.Dict[str, _typing.Any] + ] = [] self._valid_result: torch.Tensor = torch.zeros(0) self._valid_result_prob: torch.Tensor = torch.zeros(0) self._valid_score: _typing.Sequence[float] = () """ Set hyper parameters """ - self.__sampled_node_sizes: _typing.Sequence[int] = kwargs.get("sampled_node_sizes") + self.__sampled_node_sizes: _typing.Sequence[int] = kwargs.get( + "sampled_node_sizes" + ) self.__training_batch_size: int = kwargs.get("training_batch_size", 1024) if not self.__training_batch_size > 0: @@ -795,7 +831,7 @@ class NodeClassificationLayerDependentImportanceSamplingTrainer(BaseNodeClassifi optimizer: torch.optim.Optimizer = self._optimizer_class( self.model.model.parameters(), lr=self._learning_rate, - weight_decay=self._weight_decay + weight_decay=self._weight_decay, ) if type(self._lr_scheduler_type) == str: @@ -828,9 +864,11 @@ class NodeClassificationLayerDependentImportanceSamplingTrainer(BaseNodeClassifi __layer_dependent_importance_sampler: LayerDependentImportanceSampler = ( LayerDependentImportanceSampler( - integral_data.edge_index, torch.where(integral_data.train_mask)[0].unique(), - self.__sampled_node_sizes, batch_size=self.__training_batch_size, - num_workers=self.__training_sampler_num_workers + integral_data.edge_index, + torch.where(integral_data.train_mask)[0].unique(), + self.__sampled_node_sizes, + batch_size=self.__training_batch_size, + num_workers=self.__training_sampler_num_workers, ) ) for current_epoch in range(self._max_epoch): @@ -844,7 +882,7 @@ class NodeClassificationLayerDependentImportanceSamplingTrainer(BaseNodeClassifi # 由于现在的Model设计是接受Data的，所以只能组装一个采样的Data作为参数 sampled_graph: autogl.data.Data = autogl.data.Data( x=integral_data.x[sampled_data.all_sampled_nodes_indexes], - y=integral_data.y[sampled_data.all_sampled_nodes_indexes] + y=integral_data.y[sampled_data.all_sampled_nodes_indexes], ) sampled_graph.to(self.device) sampled_graph.edge_indexes: _typing.Sequence[torch.LongTensor] = [ @@ -856,17 +894,21 @@ class NodeClassificationLayerDependentImportanceSamplingTrainer(BaseNodeClassifi for current_layer in sampled_data.sampled_edges_for_layers ] if isinstance(self.model.model, ClassificationSupportedSequentialModel): - prediction: torch.Tensor = self.model.model.cls_forward(sampled_graph) + prediction: torch.Tensor = self.model.model.cls_forward( + sampled_graph + ) else: prediction: torch.Tensor = self.model.model(sampled_graph) if not hasattr(torch.nn.functional, self.loss): - raise TypeError( - f"PyTorch does not support loss type {self.loss}" - ) + raise TypeError(f"PyTorch does not support loss type {self.loss}") loss_function = getattr(torch.nn.functional, self.loss) loss_value: torch.Tensor = loss_function( - prediction[sampled_data.target_nodes_indexes.indexes_in_sampled_graph], - sampled_graph.y[sampled_data.target_nodes_indexes.indexes_in_sampled_graph] + prediction[ + sampled_data.target_nodes_indexes.indexes_in_sampled_graph + ], + sampled_graph.y[ + sampled_data.target_nodes_indexes.indexes_in_sampled_graph + ], ) loss_value.backward() optimizer.step() @@ -875,9 +917,9 @@ class NodeClassificationLayerDependentImportanceSamplingTrainer(BaseNodeClassifi lr_scheduler.step() if ( - hasattr(integral_data, "val_mask") and - getattr(integral_data, "val_mask") is not None and - type(getattr(integral_data, "val_mask")) == torch.Tensor + hasattr(integral_data, "val_mask") + and getattr(integral_data, "val_mask") is not None + and type(getattr(integral_data, "val_mask")) == torch.Tensor ): validation_results: _typing.Sequence[float] = self.evaluate( (integral_data,), "val", [self.feval[0]] @@ -891,17 +933,16 @@ class NodeClassificationLayerDependentImportanceSamplingTrainer(BaseNodeClassifi LOGGER.debug("Early stopping at %d", current_epoch) break if ( - hasattr(integral_data, "val_mask") and - getattr(integral_data, "val_mask") is not None and - type(getattr(integral_data, "val_mask")) == torch.Tensor + hasattr(integral_data, "val_mask") + and getattr(integral_data, "val_mask") is not None + and type(getattr(integral_data, "val_mask")) == torch.Tensor ): self._early_stopping.load_checkpoint(self.model.model) def __predict_only( - self, integral_data, - mask_or_target_nodes_indexes: _typing.Union[ - torch.BoolTensor, torch.LongTensor - ] + self, + integral_data, + mask_or_target_nodes_indexes: _typing.Union[torch.BoolTensor, torch.LongTensor], ) -> torch.Tensor: """ The function of predicting on the given data. @@ -911,30 +952,41 @@ class NodeClassificationLayerDependentImportanceSamplingTrainer(BaseNodeClassifi """ self.model.model.eval() integral_data = integral_data.to(torch.device("cpu")) - mask_or_target_nodes_indexes = mask_or_target_nodes_indexes.to(torch.device("cpu")) + mask_or_target_nodes_indexes = mask_or_target_nodes_indexes.to( + torch.device("cpu") + ) if isinstance(self.model.model, ClassificationSupportedSequentialModel): - sequential_gnn_model: ClassificationSupportedSequentialModel = self.model.model + sequential_gnn_model: ClassificationSupportedSequentialModel = ( + self.model.model + ) __num_layers: int = len(self.__sampled_node_sizes) x: torch.Tensor = getattr(integral_data, "x") for _current_layer_index in range(__num_layers - 1): __next_x: _typing.Optional[torch.Tensor] = None - _optional_neighbor_sampler: _typing.Optional[NeighborSampler] = ( - self.__neighbor_sampler_store[torch.arange(x.size(0))] - ) + _optional_neighbor_sampler: _typing.Optional[ + NeighborSampler + ] = self.__neighbor_sampler_store[torch.arange(x.size(0))] if ( - _optional_neighbor_sampler is not None and - type(_optional_neighbor_sampler) == NeighborSampler + _optional_neighbor_sampler is not None + and type(_optional_neighbor_sampler) == NeighborSampler ): - current_neighbor_sampler: NeighborSampler = _optional_neighbor_sampler + current_neighbor_sampler: NeighborSampler = ( + _optional_neighbor_sampler + ) else: current_neighbor_sampler: NeighborSampler = NeighborSampler( - integral_data.edge_index, torch.arange(x.size(0)).unique(), - [-1], batch_size=self.__predicting_batch_size, - num_workers=self.__predicting_sampler_num_workers, shuffle=False + integral_data.edge_index, + torch.arange(x.size(0)).unique(), + [-1], + batch_size=self.__predicting_batch_size, + num_workers=self.__predicting_sampler_num_workers, + shuffle=False, ) - self.__neighbor_sampler_store[torch.arange(x.size(0))] = current_neighbor_sampler + self.__neighbor_sampler_store[ + torch.arange(x.size(0)) + ] = current_neighbor_sampler for _target_dependant_sampled_data in current_neighbor_sampler: _target_dependant_sampled_data: TargetDependantSampledData = ( @@ -943,50 +995,63 @@ class NodeClassificationLayerDependentImportanceSamplingTrainer(BaseNodeClassifi _sampled_graph: autogl.data.Data = autogl.data.Data( x=x[_target_dependant_sampled_data.all_sampled_nodes_indexes], edge_index=( - _target_dependant_sampled_data.sampled_edges_for_layers[0].edge_index_for_sampled_graph - ) + _target_dependant_sampled_data.sampled_edges_for_layers[ + 0 + ].edge_index_for_sampled_graph + ), ) _sampled_graph.edge_weight: torch.Tensor = ( - _target_dependant_sampled_data.sampled_edges_for_layers[0].edge_weight + _target_dependant_sampled_data.sampled_edges_for_layers[ + 0 + ].edge_weight ) _sampled_graph: autogl.data.Data = _sampled_graph.to(self.device) with torch.no_grad(): __sampled_graph_inferences: torch.Tensor = ( - sequential_gnn_model.sequential_encoding_layers[_current_layer_index](_sampled_graph) + sequential_gnn_model.sequential_encoding_layers[ + _current_layer_index + ](_sampled_graph) ) _sampled_target_nodes_inferences: torch.Tensor = __sampled_graph_inferences[ _target_dependant_sampled_data.target_nodes_indexes.indexes_in_sampled_graph ].cpu() if __next_x is None: - __next_x: torch.Tensor = torch.zeros(x.size(0), __sampled_graph_inferences.size(1)) - __next_x[_target_dependant_sampled_data.target_nodes_indexes.indexes_in_integral_graph] = ( - _sampled_target_nodes_inferences - ) + __next_x: torch.Tensor = torch.zeros( + x.size(0), __sampled_graph_inferences.size(1) + ) + __next_x[ + _target_dependant_sampled_data.target_nodes_indexes.indexes_in_integral_graph + ] = _sampled_target_nodes_inferences x: torch.Tensor = __next_x " The following procedures are for the top layer " if mask_or_target_nodes_indexes.dtype == torch.bool: - target_nodes_indexes: _typing.Any = ( - torch.where(mask_or_target_nodes_indexes)[0] - ) + target_nodes_indexes: _typing.Any = torch.where( + mask_or_target_nodes_indexes + )[0] else: target_nodes_indexes: _typing.Any = mask_or_target_nodes_indexes.long() - _optional_neighbor_sampler: _typing.Optional[NeighborSampler] = ( - self.__neighbor_sampler_store[target_nodes_indexes] - ) + _optional_neighbor_sampler: _typing.Optional[ + NeighborSampler + ] = self.__neighbor_sampler_store[target_nodes_indexes] if ( - _optional_neighbor_sampler is not None and - type(_optional_neighbor_sampler) == NeighborSampler + _optional_neighbor_sampler is not None + and type(_optional_neighbor_sampler) == NeighborSampler ): current_neighbor_sampler: NeighborSampler = _optional_neighbor_sampler else: current_neighbor_sampler: NeighborSampler = NeighborSampler( - integral_data.edge_index, target_nodes_indexes, - [-1], batch_size=self.__predicting_batch_size, - num_workers=self.__predicting_sampler_num_workers, shuffle=False + integral_data.edge_index, + target_nodes_indexes, + [-1], + batch_size=self.__predicting_batch_size, + num_workers=self.__predicting_sampler_num_workers, + shuffle=False, ) - self.__neighbor_sampler_store[target_nodes_indexes] = current_neighbor_sampler + self.__neighbor_sampler_store[ + target_nodes_indexes + ] = current_neighbor_sampler prediction_batch_cumulative_builder = ( EvaluatorUtility.PredictionBatchCumulativeBuilder() @@ -994,33 +1059,46 @@ class NodeClassificationLayerDependentImportanceSamplingTrainer(BaseNodeClassifi for _target_dependant_sampled_data in current_neighbor_sampler: _sampled_graph: autogl.data.Data = autogl.data.Data( x[_target_dependant_sampled_data.all_sampled_nodes_indexes], - _target_dependant_sampled_data.sampled_edges_for_layers[0].edge_index_for_sampled_graph + _target_dependant_sampled_data.sampled_edges_for_layers[ + 0 + ].edge_index_for_sampled_graph, ) _sampled_graph.edge_weight: torch.Tensor = ( - _target_dependant_sampled_data.sampled_edges_for_layers[0].edge_weight + _target_dependant_sampled_data.sampled_edges_for_layers[ + 0 + ].edge_weight ) _sampled_graph: autogl.data.Data = _sampled_graph.to(self.device) with torch.no_grad(): prediction_batch_cumulative_builder.add_batch( _target_dependant_sampled_data.target_nodes_indexes.indexes_in_integral_graph.cpu().numpy(), sequential_gnn_model.cls_decode( - sequential_gnn_model.sequential_encoding_layers[-1](_sampled_graph) - )[_target_dependant_sampled_data.target_nodes_indexes.indexes_in_sampled_graph].cpu().numpy() + sequential_gnn_model.sequential_encoding_layers[-1]( + _sampled_graph + ) + )[ + _target_dependant_sampled_data.target_nodes_indexes.indexes_in_sampled_graph + ] + .cpu() + .numpy(), ) return torch.from_numpy(prediction_batch_cumulative_builder.compose()[1]) else: if mask_or_target_nodes_indexes.dtype == torch.bool: - target_nodes_indexes: _typing.Any = ( - torch.where(mask_or_target_nodes_indexes)[0] - ) + target_nodes_indexes: _typing.Any = torch.where( + mask_or_target_nodes_indexes + )[0] else: target_nodes_indexes: _typing.Any = mask_or_target_nodes_indexes.long() neighbor_sampler: NeighborSampler = NeighborSampler( - torch_geometric.utils.add_remaining_self_loops(integral_data.edge_index)[0], - target_nodes_indexes, [-1 for _ in self.__sampled_node_sizes], + torch_geometric.utils.add_remaining_self_loops( + integral_data.edge_index + )[0], + target_nodes_indexes, + [-1 for _ in self.__sampled_node_sizes], batch_size=self.__predicting_batch_size, num_workers=self.__predicting_sampler_num_workers, - shuffle=False + shuffle=False, ) prediction_batch_cumulative_builder = ( EvaluatorUtility.PredictionBatchCumulativeBuilder() @@ -1030,7 +1108,7 @@ class NodeClassificationLayerDependentImportanceSamplingTrainer(BaseNodeClassifi sampled_data: TargetDependantSampledData = sampled_data sampled_graph: autogl.data.Data = autogl.data.Data( integral_data.x[sampled_data.all_sampled_nodes_indexes], - integral_data.y[sampled_data.all_sampled_nodes_indexes] + integral_data.y[sampled_data.all_sampled_nodes_indexes], ) sampled_graph.to(self.device) sampled_graph.edge_indexes: _typing.Sequence[torch.LongTensor] = [ @@ -1046,13 +1124,14 @@ class NodeClassificationLayerDependentImportanceSamplingTrainer(BaseNodeClassifi sampled_data.target_nodes_indexes.indexes_in_integral_graph.cpu().numpy(), self.model.model(sampled_graph)[ sampled_data.target_nodes_indexes.indexes_in_sampled_graph - ].cpu().numpy() + ] + .cpu() + .numpy(), ) return torch.from_numpy(prediction_batch_cumulative_builder.compose()[1]) def predict_proba( - self, dataset, mask: _typing.Optional[str] = None, - in_log_format: bool = False + self, dataset, mask: _typing.Optional[str] = None, in_log_format: bool = False ): """ The function of predicting the probability on the given dataset. @@ -1107,12 +1186,12 @@ class NodeClassificationLayerDependentImportanceSamplingTrainer(BaseNodeClassifi return self.predict_proba(dataset, mask, in_log_format=True).max(1)[1] def evaluate( - self, - dataset, - mask: _typing.Optional[str] = None, - feval: _typing.Union[ - None, _typing.Sequence[str], _typing.Sequence[_typing.Type[Evaluation]] - ] = None, + self, + dataset, + mask: _typing.Optional[str] = None, + feval: _typing.Union[ + None, _typing.Sequence[str], _typing.Sequence[_typing.Type[Evaluation]] + ] = None, ) -> _typing.Sequence[float]: """ The function of training on the given dataset and keeping valid result. @@ -1156,7 +1235,8 @@ class NodeClassificationLayerDependentImportanceSamplingTrainer(BaseNodeClassifi f.evaluate( prediction_probability.cpu().numpy(), y_ground_truth.cpu().numpy(), - ) for f in _feval + ) + for f in _feval ] def train(self, dataset, keep_valid_result: bool = True): @@ -1176,6 +1256,7 @@ class NodeClassificationLayerDependentImportanceSamplingTrainer(BaseNodeClassifi None """ import gc + gc.collect() data = dataset[0] self.__train_only(data) @@ -1193,10 +1274,10 @@ class NodeClassificationLayerDependentImportanceSamplingTrainer(BaseNodeClassifi return self._valid_result_prob def get_valid_score( - self, return_major: bool = True + self, return_major: bool = True ) -> _typing.Union[ _typing.Tuple[float, bool], - _typing.Tuple[_typing.Sequence[float], _typing.Sequence[bool]] + _typing.Tuple[_typing.Sequence[float], _typing.Sequence[bool]], ]: """ The function of getting the valid score. @@ -1222,7 +1303,7 @@ class NodeClassificationLayerDependentImportanceSamplingTrainer(BaseNodeClassifi @hyper_parameter_space.setter def hyper_parameter_space( - self, hp_space: _typing.Sequence[_typing.Dict[str, _typing.Any]] + self, hp_space: _typing.Sequence[_typing.Dict[str, _typing.Any]] ) -> None: if not isinstance(hp_space, _typing.Sequence): raise TypeError @@ -1230,6 +1311,7 @@ class NodeClassificationLayerDependentImportanceSamplingTrainer(BaseNodeClassifi def __repr__(self) -> str: import yaml + return yaml.dump( { "trainer_name": self.__class__.__name__, @@ -1238,14 +1320,14 @@ class NodeClassificationLayerDependentImportanceSamplingTrainer(BaseNodeClassifi "max_epoch": self._max_epoch, "early_stopping_round": self._early_stopping.patience, "sampling_sizes": self.__sampled_node_sizes, - "model": repr(self.model) + "model": repr(self.model), } ) def duplicate_from_hyper_parameter( - self, - hp: _typing.Dict[str, _typing.Any], - model: _typing.Optional[BaseModel] = None, + self, + hp: _typing.Dict[str, _typing.Any], + model: _typing.Optional[BaseModel] = None, ) -> "NodeClassificationLayerDependentImportanceSamplingTrainer": """ The function of duplicating a new instance from the given hyper-parameter. @@ -1317,24 +1399,25 @@ class NodeClassificationNeighborSamplingTrainer(BaseNodeClassificationTrainer): feval: ``str``. The evaluation method adopted in this function. """ + def __init__( - self, - model: _typing.Union[BaseModel, str], - num_features: int, - num_classes: int, - optimizer: _typing.Union[_typing.Type[torch.optim.Optimizer], str, None] = ..., - lr: float = 1e-4, - max_epoch: int = 100, - early_stopping_round: int = 100, - weight_decay: float = 1e-4, - device: _typing.Optional[torch.device] = None, - init: bool = True, - feval: _typing.Union[ - _typing.Sequence[str], _typing.Sequence[_typing.Type[Evaluation]] - ] = (Logloss,), - loss: str = "nll_loss", - lr_scheduler_type: _typing.Optional[str] = None, - **kwargs, + self, + model: _typing.Union[BaseModel, str], + num_features: int, + num_classes: int, + optimizer: _typing.Union[_typing.Type[torch.optim.Optimizer], str, None] = ..., + lr: float = 1e-4, + max_epoch: int = 100, + early_stopping_round: int = 100, + weight_decay: float = 1e-4, + device: _typing.Optional[torch.device] = None, + init: bool = True, + feval: _typing.Union[ + _typing.Sequence[str], _typing.Sequence[_typing.Type[Evaluation]] + ] = (Logloss,), + loss: str = "nll_loss", + lr_scheduler_type: _typing.Optional[str] = None, + **kwargs, ): if isinstance(optimizer, type) and issubclass(optimizer, torch.optim.Optimizer): self._optimizer_class: _typing.Type[torch.optim.Optimizer] = optimizer @@ -1365,10 +1448,12 @@ class NodeClassificationNeighborSamplingTrainer(BaseNodeClassificationTrainer): self._weight_decay: float = weight_decay if weight_decay > 0 else 1e-4 self._early_stopping = EarlyStopping( patience=early_stopping_round if early_stopping_round > 0 else 1e2, - verbose=False + verbose=False, ) """ Assign an empty initial hyper parameter space """ - self._hyper_parameter_space: _typing.Sequence[_typing.Dict[str, _typing.Any]] = [] + self._hyper_parameter_space: _typing.Sequence[ + _typing.Dict[str, _typing.Any] + ] = [] self._valid_result: torch.Tensor = torch.zeros(0) self._valid_result_prob: torch.Tensor = torch.zeros(0) @@ -1432,7 +1517,7 @@ class NodeClassificationNeighborSamplingTrainer(BaseNodeClassificationTrainer): optimizer: torch.optim.Optimizer = self._optimizer_class( self.model.model.parameters(), lr=self._learning_rate, - weight_decay=self._weight_decay + weight_decay=self._weight_decay, ) if type(self._lr_scheduler_type) == str: @@ -1464,9 +1549,11 @@ class NodeClassificationNeighborSamplingTrainer(BaseNodeClassificationTrainer): ) neighbor_sampler: NeighborSampler = NeighborSampler( - integral_data.edge_index, torch.where(integral_data.train_mask)[0].unique(), - self.__sampling_sizes, batch_size=self.__training_batch_size, - num_workers=self.__training_sampler_num_workers + integral_data.edge_index, + torch.where(integral_data.train_mask)[0].unique(), + self.__sampling_sizes, + batch_size=self.__training_batch_size, + num_workers=self.__training_sampler_num_workers, ) for current_epoch in range(self._max_epoch): self.model.model.train() @@ -1478,7 +1565,7 @@ class NodeClassificationNeighborSamplingTrainer(BaseNodeClassificationTrainer): sampled_data: TargetDependantSampledData = sampled_data sampled_graph: autogl.data.Data = autogl.data.Data( x=integral_data.x[sampled_data.all_sampled_nodes_indexes], - y=integral_data.y[sampled_data.all_sampled_nodes_indexes] + y=integral_data.y[sampled_data.all_sampled_nodes_indexes], ) sampled_graph.to(self.device) sampled_graph.edge_indexes: _typing.Sequence[torch.LongTensor] = [ @@ -1486,17 +1573,21 @@ class NodeClassificationNeighborSamplingTrainer(BaseNodeClassificationTrainer): for current_layer in sampled_data.sampled_edges_for_layers ] if isinstance(self.model.model, ClassificationSupportedSequentialModel): - prediction: torch.Tensor = self.model.model.cls_forward(sampled_graph) + prediction: torch.Tensor = self.model.model.cls_forward( + sampled_graph + ) else: prediction: torch.Tensor = self.model.model(sampled_graph) if not hasattr(torch.nn.functional, self.loss): - raise TypeError( - f"PyTorch does not support loss type {self.loss}" - ) + raise TypeError(f"PyTorch does not support loss type {self.loss}") loss_function = getattr(torch.nn.functional, self.loss) loss_value: torch.Tensor = loss_function( - prediction[sampled_data.target_nodes_indexes.indexes_in_sampled_graph], - sampled_graph.y[sampled_data.target_nodes_indexes.indexes_in_sampled_graph] + prediction[ + sampled_data.target_nodes_indexes.indexes_in_sampled_graph + ], + sampled_graph.y[ + sampled_data.target_nodes_indexes.indexes_in_sampled_graph + ], ) loss_value.backward() optimizer.step() @@ -1505,9 +1596,9 @@ class NodeClassificationNeighborSamplingTrainer(BaseNodeClassificationTrainer): lr_scheduler.step() if ( - hasattr(integral_data, "val_mask") and - getattr(integral_data, "val_mask") is not None and - type(getattr(integral_data, "val_mask")) == torch.Tensor + hasattr(integral_data, "val_mask") + and getattr(integral_data, "val_mask") is not None + and type(getattr(integral_data, "val_mask")) == torch.Tensor ): validation_results: _typing.Sequence[float] = self.evaluate( (integral_data,), "val", [self.feval[0]] @@ -1521,17 +1612,16 @@ class NodeClassificationNeighborSamplingTrainer(BaseNodeClassificationTrainer): LOGGER.debug("Early stopping at %d", current_epoch) break if ( - hasattr(integral_data, "val_mask") and - getattr(integral_data, "val_mask") is not None and - type(getattr(integral_data, "val_mask")) == torch.Tensor + hasattr(integral_data, "val_mask") + and getattr(integral_data, "val_mask") is not None + and type(getattr(integral_data, "val_mask")) == torch.Tensor ): self._early_stopping.load_checkpoint(self.model.model) def __predict_only( - self, integral_data, - mask_or_target_nodes_indexes: _typing.Union[ - torch.BoolTensor, torch.LongTensor - ] + self, + integral_data, + mask_or_target_nodes_indexes: _typing.Union[torch.BoolTensor, torch.LongTensor], ) -> torch.Tensor: """ The function of predicting on the given data. @@ -1541,28 +1631,37 @@ class NodeClassificationNeighborSamplingTrainer(BaseNodeClassificationTrainer): """ self.model.model.eval() integral_data = integral_data.to(torch.device("cpu")) - mask_or_target_nodes_indexes = mask_or_target_nodes_indexes.to(torch.device("cpu")) + mask_or_target_nodes_indexes = mask_or_target_nodes_indexes.to( + torch.device("cpu") + ) if isinstance(self.model.model, ClassificationSupportedSequentialModel): - sequential_gnn_model: ClassificationSupportedSequentialModel = self.model.model + sequential_gnn_model: ClassificationSupportedSequentialModel = ( + self.model.model + ) __num_layers: int = len(self.__sampling_sizes) x: torch.Tensor = getattr(integral_data, "x") for _current_layer_index in range(__num_layers - 1): __next_x: _typing.Optional[torch.Tensor] = None - _optional_neighbor_sampler: _typing.Optional[NeighborSampler] = ( - self.__neighbor_sampler_store[torch.arange(x.size(0)).unique()] - ) + _optional_neighbor_sampler: _typing.Optional[ + NeighborSampler + ] = self.__neighbor_sampler_store[torch.arange(x.size(0)).unique()] if ( - _optional_neighbor_sampler is not None and - type(_optional_neighbor_sampler) == NeighborSampler + _optional_neighbor_sampler is not None + and type(_optional_neighbor_sampler) == NeighborSampler ): - current_neighbor_sampler: NeighborSampler = _optional_neighbor_sampler + current_neighbor_sampler: NeighborSampler = ( + _optional_neighbor_sampler + ) else: current_neighbor_sampler: NeighborSampler = NeighborSampler( - integral_data.edge_index, torch.arange(x.size(0)).unique(), - [-1], batch_size=self.__predicting_batch_size, - num_workers=self.__predicting_sampler_num_workers, shuffle=False + integral_data.edge_index, + torch.arange(x.size(0)).unique(), + [-1], + batch_size=self.__predicting_batch_size, + num_workers=self.__predicting_sampler_num_workers, + shuffle=False, ) __temp: _typing.Any = torch.arange(x.size(0)) self.__neighbor_sampler_store[__temp] = current_neighbor_sampler @@ -1574,47 +1673,58 @@ class NodeClassificationNeighborSamplingTrainer(BaseNodeClassificationTrainer): _sampled_graph: autogl.data.Data = autogl.data.Data( x=x[_target_dependant_sampled_data.all_sampled_nodes_indexes], edge_index=( - _target_dependant_sampled_data.sampled_edges_for_layers[0].edge_index_for_sampled_graph - ) + _target_dependant_sampled_data.sampled_edges_for_layers[ + 0 + ].edge_index_for_sampled_graph + ), ) _sampled_graph: autogl.data.Data = _sampled_graph.to(self.device) with torch.no_grad(): __sampled_graph_inferences: torch.Tensor = ( - sequential_gnn_model.sequential_encoding_layers[_current_layer_index](_sampled_graph) + sequential_gnn_model.sequential_encoding_layers[ + _current_layer_index + ](_sampled_graph) ) _sampled_target_nodes_inferences: torch.Tensor = __sampled_graph_inferences[ _target_dependant_sampled_data.target_nodes_indexes.indexes_in_sampled_graph ].cpu() if __next_x is None: - __next_x: torch.Tensor = torch.zeros(x.size(0), __sampled_graph_inferences.size(1)) - __next_x[_target_dependant_sampled_data.target_nodes_indexes.indexes_in_integral_graph] = ( - _sampled_target_nodes_inferences - ) + __next_x: torch.Tensor = torch.zeros( + x.size(0), __sampled_graph_inferences.size(1) + ) + __next_x[ + _target_dependant_sampled_data.target_nodes_indexes.indexes_in_integral_graph + ] = _sampled_target_nodes_inferences x: torch.Tensor = __next_x # The following procedures are for the top layer if mask_or_target_nodes_indexes.dtype == torch.bool: - target_nodes_indexes: _typing.Any = ( - torch.where(mask_or_target_nodes_indexes)[0] - ) + target_nodes_indexes: _typing.Any = torch.where( + mask_or_target_nodes_indexes + )[0] else: target_nodes_indexes: _typing.Any = mask_or_target_nodes_indexes.long() - _optional_neighbor_sampler: _typing.Optional[NeighborSampler] = ( - self.__neighbor_sampler_store[target_nodes_indexes] - ) + _optional_neighbor_sampler: _typing.Optional[ + NeighborSampler + ] = self.__neighbor_sampler_store[target_nodes_indexes] if ( - _optional_neighbor_sampler is not None and - type(_optional_neighbor_sampler) == NeighborSampler + _optional_neighbor_sampler is not None + and type(_optional_neighbor_sampler) == NeighborSampler ): current_neighbor_sampler: NeighborSampler = _optional_neighbor_sampler else: current_neighbor_sampler: NeighborSampler = NeighborSampler( - integral_data.edge_index, target_nodes_indexes, - [-1], batch_size=self.__predicting_batch_size, - num_workers=self.__predicting_sampler_num_workers, shuffle=False + integral_data.edge_index, + target_nodes_indexes, + [-1], + batch_size=self.__predicting_batch_size, + num_workers=self.__predicting_sampler_num_workers, + shuffle=False, ) - self.__neighbor_sampler_store[target_nodes_indexes] = current_neighbor_sampler + self.__neighbor_sampler_store[ + target_nodes_indexes + ] = current_neighbor_sampler prediction_batch_cumulative_builder = ( EvaluatorUtility.PredictionBatchCumulativeBuilder() @@ -1622,29 +1732,40 @@ class NodeClassificationNeighborSamplingTrainer(BaseNodeClassificationTrainer): for _target_dependant_sampled_data in current_neighbor_sampler: _sampled_graph: autogl.data.Data = autogl.data.Data( x[_target_dependant_sampled_data.all_sampled_nodes_indexes], - _target_dependant_sampled_data.sampled_edges_for_layers[0].edge_index_for_sampled_graph + _target_dependant_sampled_data.sampled_edges_for_layers[ + 0 + ].edge_index_for_sampled_graph, ) _sampled_graph: autogl.data.Data = _sampled_graph.to(self.device) with torch.no_grad(): prediction_batch_cumulative_builder.add_batch( _target_dependant_sampled_data.target_nodes_indexes.indexes_in_integral_graph.cpu().numpy(), sequential_gnn_model.cls_decode( - sequential_gnn_model.sequential_encoding_layers[-1](_sampled_graph) - )[_target_dependant_sampled_data.target_nodes_indexes.indexes_in_sampled_graph].cpu().numpy() + sequential_gnn_model.sequential_encoding_layers[-1]( + _sampled_graph + ) + )[ + _target_dependant_sampled_data.target_nodes_indexes.indexes_in_sampled_graph + ] + .cpu() + .numpy(), ) return torch.from_numpy(prediction_batch_cumulative_builder.compose()[1]) else: if mask_or_target_nodes_indexes.dtype == torch.bool: - target_nodes_indexes: _typing.Any = ( - torch.where(mask_or_target_nodes_indexes)[0] - ) + target_nodes_indexes: _typing.Any = torch.where( + mask_or_target_nodes_indexes + )[0] else: target_nodes_indexes: _typing.Any = mask_or_target_nodes_indexes.long() neighbor_sampler: NeighborSampler = NeighborSampler( - integral_data.edge_index, target_nodes_indexes, [-1 for _ in self.__sampling_sizes], + integral_data.edge_index, + target_nodes_indexes, + [-1 for _ in self.__sampling_sizes], batch_size=self.__predicting_batch_size, - num_workers=self.__predicting_sampler_num_workers, shuffle=False + num_workers=self.__predicting_sampler_num_workers, + shuffle=False, ) prediction_batch_cumulative_builder = ( @@ -1667,13 +1788,14 @@ class NodeClassificationNeighborSamplingTrainer(BaseNodeClassificationTrainer): _target_dependant_sampled_data.target_nodes_indexes.indexes_in_integral_graph.cpu().numpy(), self.model.model(_sampled_graph)[ _target_dependant_sampled_data.target_nodes_indexes.indexes_in_sampled_graph - ].cpu().numpy() + ] + .cpu() + .numpy(), ) return torch.from_numpy(prediction_batch_cumulative_builder.compose()[1]) def predict_proba( - self, dataset, mask: _typing.Optional[str] = None, - in_log_format: bool = False + self, dataset, mask: _typing.Optional[str] = None, in_log_format: bool = False ): """ The function of predicting the probability on the given dataset. @@ -1728,12 +1850,12 @@ class NodeClassificationNeighborSamplingTrainer(BaseNodeClassificationTrainer): return self.predict_proba(dataset, mask, in_log_format=True).max(1)[1] def evaluate( - self, - dataset, - mask: _typing.Optional[str] = None, - feval: _typing.Union[ - None, _typing.Sequence[str], _typing.Sequence[_typing.Type[Evaluation]] - ] = None, + self, + dataset, + mask: _typing.Optional[str] = None, + feval: _typing.Union[ + None, _typing.Sequence[str], _typing.Sequence[_typing.Type[Evaluation]] + ] = None, ) -> _typing.Sequence[float]: """ The function of training on the given dataset and keeping valid result. @@ -1800,6 +1922,7 @@ class NodeClassificationNeighborSamplingTrainer(BaseNodeClassificationTrainer): None """ import gc + gc.collect() data = dataset[0] self.__train_only(data) @@ -1817,10 +1940,10 @@ class NodeClassificationNeighborSamplingTrainer(BaseNodeClassificationTrainer): return self._valid_result_prob def get_valid_score( - self, return_major: bool = True + self, return_major: bool = True ) -> _typing.Union[ _typing.Tuple[float, bool], - _typing.Tuple[_typing.Sequence[float], _typing.Sequence[bool]] + _typing.Tuple[_typing.Sequence[float], _typing.Sequence[bool]], ]: """ The function of getting the valid score. @@ -1846,7 +1969,7 @@ class NodeClassificationNeighborSamplingTrainer(BaseNodeClassificationTrainer): @hyper_parameter_space.setter def hyper_parameter_space( - self, hp_space: _typing.Sequence[_typing.Dict[str, _typing.Any]] + self, hp_space: _typing.Sequence[_typing.Dict[str, _typing.Any]] ) -> None: if not isinstance(hp_space, _typing.Sequence): raise TypeError @@ -1854,6 +1977,7 @@ class NodeClassificationNeighborSamplingTrainer(BaseNodeClassificationTrainer): def __repr__(self) -> str: import yaml + return yaml.dump( { "trainer_name": self.__class__.__name__, @@ -1862,14 +1986,14 @@ class NodeClassificationNeighborSamplingTrainer(BaseNodeClassificationTrainer): "max_epoch": self._max_epoch, "early_stopping_round": self._early_stopping.patience, "sampling_sizes": self.__sampling_sizes, - "model": repr(self.model) + "model": repr(self.model), } ) def duplicate_from_hyper_parameter( - self, - hp: _typing.Dict[str, _typing.Any], - model: _typing.Optional[BaseModel] = None, + self, + hp: _typing.Dict[str, _typing.Any], + model: _typing.Optional[BaseModel] = None, ) -> "NodeClassificationNeighborSamplingTrainer": """ The function of duplicating a new instance from the given hyper-parameter. diff --git a/autogl/module/train/sampling/sampler/graphsaint_sampler.py b/autogl/module/train/sampling/sampler/graphsaint_sampler.py index 8948990..00dbded 100644 --- a/autogl/module/train/sampling/sampler/graphsaint_sampler.py +++ b/autogl/module/train/sampling/sampler/graphsaint_sampler.py @@ -19,10 +19,15 @@ class GraphSAINTSamplerFactory: With the aim of abstracting a unified sampling module for representative mainstream varieties of Node-wise Sampling, Layer-wise Sampling, and Subgraph-wise Sampling. """ + @classmethod def create_node_sampler( - cls, data, num_graphs_per_epoch: int, node_budget: int, - sample_coverage_factor: int = 50, **kwargs + cls, + data, + num_graphs_per_epoch: int, + node_budget: int, + sample_coverage_factor: int = 50, + **kwargs ) -> torch_geometric.data.GraphSAINTNodeSampler: """ A simple static method for instantiating :class:`torch_geometric.data.GraphSAINTNodeSampler` @@ -48,14 +53,22 @@ class GraphSAINTSamplerFactory: Instance of :class:`torch_geometric.data.GraphSAINTNodeSampler`. """ return torch_geometric.data.GraphSAINTNodeSampler( - data, node_budget, - num_graphs_per_epoch, sample_coverage_factor, log=False, **kwargs + data, + node_budget, + num_graphs_per_epoch, + sample_coverage_factor, + log=False, + **kwargs ) @classmethod def create_edge_sampler( - cls, data, num_graphs_per_epoch: int, edge_budget: int, - sample_coverage_factor: int = 50, **kwargs + cls, + data, + num_graphs_per_epoch: int, + edge_budget: int, + sample_coverage_factor: int = 50, + **kwargs ) -> torch_geometric.data.GraphSAINTEdgeSampler: """ A simple static method for instantiating :class:`torch_geometric.data.GraphSAINTEdgeSampler` @@ -81,15 +94,23 @@ class GraphSAINTSamplerFactory: Instance of :class:`torch_geometric.data.GraphSAINTEdgeSampler`. """ return torch_geometric.data.GraphSAINTEdgeSampler( - data, edge_budget, - num_graphs_per_epoch, sample_coverage_factor, log=False, **kwargs + data, + edge_budget, + num_graphs_per_epoch, + sample_coverage_factor, + log=False, + **kwargs ) @classmethod def create_random_walk_sampler( - cls, data, num_graphs_per_epoch: int, - num_walks: int, walk_length: int, - sample_coverage_factor: int = 50, **kwargs + cls, + data, + num_graphs_per_epoch: int, + num_walks: int, + walk_length: int, + sample_coverage_factor: int = 50, + **kwargs ) -> torch_geometric.data.GraphSAINTRandomWalkSampler: """ A simple static method for instantiating :class:`torch_geometric.data.GraphSAINTEdgeSampler` @@ -117,6 +138,11 @@ class GraphSAINTSamplerFactory: Instance of :class:`torch_geometric.data.GraphSAINTEdgeSampler`. """ return torch_geometric.data.GraphSAINTRandomWalkSampler( - data, num_walks, walk_length, - num_graphs_per_epoch, sample_coverage_factor, log=False, **kwargs + data, + num_walks, + walk_length, + num_graphs_per_epoch, + sample_coverage_factor, + log=False, + **kwargs ) diff --git a/autogl/module/train/sampling/sampler/layer_dependent_importance_sampler.py b/autogl/module/train/sampling/sampler/layer_dependent_importance_sampler.py index ffb7df5..bc66bdf 100644 --- a/autogl/module/train/sampling/sampler/layer_dependent_importance_sampler.py +++ b/autogl/module/train/sampling/sampler/layer_dependent_importance_sampler.py @@ -7,22 +7,29 @@ import torch_geometric from . import target_dependant_sampler -class _LayerDependentImportanceSampler(target_dependant_sampler.BasicLayerWiseTargetDependantSampler): +class _LayerDependentImportanceSampler( + target_dependant_sampler.BasicLayerWiseTargetDependantSampler +): """ Obsolete implementation, unused """ + class _Utility: @classmethod - def compute_edge_weights(cls, __all_edge_index_with_self_loops: torch.Tensor) -> torch.Tensor: - __out_degree: torch.Tensor = \ - torch_geometric.utils.degree(__all_edge_index_with_self_loops[0]) - __in_degree: torch.Tensor = \ - torch_geometric.utils.degree(__all_edge_index_with_self_loops[1]) + def compute_edge_weights( + cls, __all_edge_index_with_self_loops: torch.Tensor + ) -> torch.Tensor: + __out_degree: torch.Tensor = torch_geometric.utils.degree( + __all_edge_index_with_self_loops[0] + ) + __in_degree: torch.Tensor = torch_geometric.utils.degree( + __all_edge_index_with_self_loops[1] + ) temp_tensor: torch.Tensor = torch.stack( [ __out_degree[__all_edge_index_with_self_loops[0]], - __in_degree[__all_edge_index_with_self_loops[1]] + __in_degree[__all_edge_index_with_self_loops[1]], ] ) temp_tensor: torch.Tensor = torch.pow(temp_tensor, -0.5) @@ -31,9 +38,10 @@ class _LayerDependentImportanceSampler(target_dependant_sampler.BasicLayerWiseTa @classmethod def get_candidate_source_nodes_probabilities( - cls, all_candidate_edge_indexes: torch.LongTensor, - all_edge_index_with_self_loops: torch.Tensor, - all_edge_weights: torch.Tensor + cls, + all_candidate_edge_indexes: torch.LongTensor, + all_edge_index_with_self_loops: torch.Tensor, + all_edge_weights: torch.Tensor, ) -> _typing.Tuple[torch.LongTensor, torch.Tensor]: """ :param all_candidate_edge_indexes: @@ -41,8 +49,12 @@ class _LayerDependentImportanceSampler(target_dependant_sampler.BasicLayerWiseTa :param all_edge_weights: :return: (all_source_nodes_indexes, all_source_nodes_probabilities) """ - all_candidate_edge_indexes: torch.LongTensor = all_candidate_edge_indexes.unique() - _all_candidate_edges_weights: torch.Tensor = all_edge_weights[all_candidate_edge_indexes] + all_candidate_edge_indexes: torch.LongTensor = ( + all_candidate_edge_indexes.unique() + ) + _all_candidate_edges_weights: torch.Tensor = all_edge_weights[ + all_candidate_edge_indexes + ] all_candidate_source_nodes_indexes: torch.LongTensor = ( all_edge_index_with_self_loops[0, all_candidate_edge_indexes].unique() ) @@ -50,23 +62,31 @@ class _LayerDependentImportanceSampler(target_dependant_sampler.BasicLayerWiseTa [ torch.sum( _all_candidate_edges_weights[ - all_edge_index_with_self_loops[0, all_candidate_edge_indexes] == _current_source_node_index + all_edge_index_with_self_loops[ + 0, all_candidate_edge_indexes ] - ).item() / torch.sum(_all_candidate_edges_weights).item() + == _current_source_node_index + ] + ).item() + / torch.sum(_all_candidate_edges_weights).item() for _current_source_node_index in all_candidate_source_nodes_indexes.tolist() ] ) assert ( - all_candidate_source_nodes_indexes.size() == - all_candidate_source_nodes_probabilities.size() + all_candidate_source_nodes_indexes.size() + == all_candidate_source_nodes_probabilities.size() + ) + return ( + all_candidate_source_nodes_indexes, + all_candidate_source_nodes_probabilities, ) - return all_candidate_source_nodes_indexes, all_candidate_source_nodes_probabilities @classmethod def filter_selected_edges_by_source_nodes_and_target_nodes( - cls, all_edges_with_self_loops: torch.Tensor, - selected_source_node_indexes: torch.LongTensor, - selected_target_node_indexes: torch.LongTensor + cls, + all_edges_with_self_loops: torch.Tensor, + selected_source_node_indexes: torch.LongTensor, + selected_target_node_indexes: torch.LongTensor, ) -> torch.Tensor: """ :param all_edges_with_self_loops: all edges with self loops @@ -78,41 +98,65 @@ class _LayerDependentImportanceSampler(target_dependant_sampler.BasicLayerWiseTa all_edges_with_self_loops.size(1), dtype=torch.bool ) selected_edges_mask_for_source_nodes[ - torch.cat([ - torch.where(all_edges_with_self_loops[0] == __current_selected_source_node_index)[0] - for __current_selected_source_node_index in selected_source_node_indexes.unique().tolist() - ]).unique() + torch.cat( + [ + torch.where( + all_edges_with_self_loops[0] + == __current_selected_source_node_index + )[0] + for __current_selected_source_node_index in selected_source_node_indexes.unique().tolist() + ] + ).unique() ] = True selected_edges_mask_for_target_nodes: torch.Tensor = torch.zeros( all_edges_with_self_loops.size(1), dtype=torch.bool ) selected_edges_mask_for_target_nodes[ - torch.cat([ - torch.where(all_edges_with_self_loops[1] == __current_selected_target_node_index)[0] - for __current_selected_target_node_index in selected_target_node_indexes.unique().tolist() - ]) + torch.cat( + [ + torch.where( + all_edges_with_self_loops[1] + == __current_selected_target_node_index + )[0] + for __current_selected_target_node_index in selected_target_node_indexes.unique().tolist() + ] + ) ] = True return torch.where( - selected_edges_mask_for_source_nodes & selected_edges_mask_for_target_nodes + selected_edges_mask_for_source_nodes + & selected_edges_mask_for_target_nodes )[0] def __init__( - self, edge_index: torch.LongTensor, - target_nodes_indexes: torch.LongTensor, - layer_wise_arguments: _typing.Sequence, - batch_size: _typing.Optional[int] = 1, num_workers: int = 0, - shuffle: bool = True, **kwargs + self, + edge_index: torch.LongTensor, + target_nodes_indexes: torch.LongTensor, + layer_wise_arguments: _typing.Sequence, + batch_size: _typing.Optional[int] = 1, + num_workers: int = 0, + shuffle: bool = True, + **kwargs ): super().__init__( torch_geometric.utils.add_remaining_self_loops(edge_index)[0], - target_nodes_indexes, layer_wise_arguments, batch_size, num_workers, shuffle, **kwargs + target_nodes_indexes, + layer_wise_arguments, + batch_size, + num_workers, + shuffle, + **kwargs + ) + self.__all_edge_weights: torch.Tensor = self._Utility.compute_edge_weights( + self._edge_index ) - self.__all_edge_weights: torch.Tensor = self._Utility.compute_edge_weights(self._edge_index) def _sample_edges_for_layer( - self, __current_layer_target_nodes_indexes: torch.LongTensor, - __top_layer_target_nodes_indexes: torch.LongTensor, - layer_argument: _typing.Any, *args, **kwargs + self, + __current_layer_target_nodes_indexes: torch.LongTensor, + __top_layer_target_nodes_indexes: torch.LongTensor, + layer_argument: _typing.Any, + *args, + **kwargs ) -> _typing.Tuple[torch.LongTensor, _typing.Optional[torch.Tensor]]: """ Sample edges for one layer @@ -136,71 +180,90 @@ class _LayerDependentImportanceSampler(target_dependant_sampler.BasicLayerWiseTa for current_target_node_index in __current_layer_target_nodes_indexes.unique().tolist() ] ).unique() - __all_candidate_source_nodes_indexes, all_candidate_source_nodes_probabilities = \ - self._Utility.get_candidate_source_nodes_probabilities( - all_candidate_edge_indexes, self._edge_index, - self.__all_edge_weights * self.__all_edge_weights - ) - assert __all_candidate_source_nodes_indexes.size() == all_candidate_source_nodes_probabilities.size() + ( + __all_candidate_source_nodes_indexes, + all_candidate_source_nodes_probabilities, + ) = self._Utility.get_candidate_source_nodes_probabilities( + all_candidate_edge_indexes, + self._edge_index, + self.__all_edge_weights * self.__all_edge_weights, + ) + assert ( + __all_candidate_source_nodes_indexes.size() + == all_candidate_source_nodes_probabilities.size() + ) """ Sampling """ if sampled_node_size_budget < __all_candidate_source_nodes_indexes.numel(): - selected_source_node_indexes: torch.LongTensor = __all_candidate_source_nodes_indexes[ - torch.from_numpy( - np.unique(np.random.choice( - np.arange(__all_candidate_source_nodes_indexes.numel()), sampled_node_size_budget, - p=all_candidate_source_nodes_probabilities.numpy(), replace=False - )) - ).unique() - ].unique() + selected_source_node_indexes: torch.LongTensor = ( + __all_candidate_source_nodes_indexes[ + torch.from_numpy( + np.unique( + np.random.choice( + np.arange(__all_candidate_source_nodes_indexes.numel()), + sampled_node_size_budget, + p=all_candidate_source_nodes_probabilities.numpy(), + replace=False, + ) + ) + ).unique() + ].unique() + ) else: - selected_source_node_indexes: torch.LongTensor = __all_candidate_source_nodes_indexes + selected_source_node_indexes: torch.LongTensor = ( + __all_candidate_source_nodes_indexes + ) selected_source_node_indexes: torch.LongTensor = torch.cat( [selected_source_node_indexes, __top_layer_target_nodes_indexes] ).unique() __selected_edges_indexes: torch.LongTensor = ( self._Utility.filter_selected_edges_by_source_nodes_and_target_nodes( - self._edge_index, selected_source_node_indexes, __current_layer_target_nodes_indexes + self._edge_index, + selected_source_node_indexes, + __current_layer_target_nodes_indexes, ).unique() ) - non_normalized_selected_edges_weight: torch.Tensor = ( - self.__all_edge_weights[__selected_edges_indexes] / - torch.tensor( - [ - all_candidate_source_nodes_probabilities[ - __all_candidate_source_nodes_indexes == current_source_node_index - ].item() - for current_source_node_index - in self._edge_index[0, __selected_edges_indexes].tolist() - ] - ) + non_normalized_selected_edges_weight: torch.Tensor = self.__all_edge_weights[ + __selected_edges_indexes + ] / torch.tensor( + [ + all_candidate_source_nodes_probabilities[ + __all_candidate_source_nodes_indexes == current_source_node_index + ].item() + for current_source_node_index in self._edge_index[ + 0, __selected_edges_indexes + ].tolist() + ] ) def __normalize_edges_weight_by_target_nodes( - __edge_index: torch.Tensor, __edge_weight: torch.Tensor + __edge_index: torch.Tensor, __edge_weight: torch.Tensor ) -> torch.Tensor: if __edge_index.size(1) != __edge_weight.numel(): raise ValueError for current_target_node_index in __edge_index[1].unique().tolist(): __current_mask_for_edges: torch.BoolTensor = ( - __edge_index[1] == current_target_node_index - ) - __edge_weight[__current_mask_for_edges] = ( - __edge_weight[__current_mask_for_edges] / - torch.sum(__edge_weight[__current_mask_for_edges]) + __edge_index[1] == current_target_node_index ) + __edge_weight[__current_mask_for_edges] = __edge_weight[ + __current_mask_for_edges + ] / torch.sum(__edge_weight[__current_mask_for_edges]) return __edge_weight - normalized_selected_edges_weight: torch.Tensor = __normalize_edges_weight_by_target_nodes( - self._edge_index[:, __selected_edges_indexes], - non_normalized_selected_edges_weight + normalized_selected_edges_weight: torch.Tensor = ( + __normalize_edges_weight_by_target_nodes( + self._edge_index[:, __selected_edges_indexes], + non_normalized_selected_edges_weight, + ) ) return __selected_edges_indexes, normalized_selected_edges_weight -class LayerDependentImportanceSampler(target_dependant_sampler.BasicLayerWiseTargetDependantSampler): +class LayerDependentImportanceSampler( + target_dependant_sampler.BasicLayerWiseTargetDependantSampler +): """ The layer-dependent importance sampler from the `"Layer-Dependent Importance Sampling for Training Deep and Large Graph Convolutional Networks" @@ -236,38 +299,54 @@ class LayerDependentImportanceSampler(target_dependant_sampler.BasicLayerWiseTar @classmethod def __compute_edge_weight(cls, edge_index: torch.Tensor) -> torch.Tensor: __num_nodes: int = max(int(edge_index[0].max()), int(edge_index[1].max())) + 1 - _temp_tensor: torch.Tensor = torch.stack([ - torch_geometric.utils.degree(edge_index[0], __num_nodes)[edge_index[0]], - torch_geometric.utils.degree(edge_index[1], __num_nodes)[edge_index[1]] - ]) + _temp_tensor: torch.Tensor = torch.stack( + [ + torch_geometric.utils.degree(edge_index[0], __num_nodes)[edge_index[0]], + torch_geometric.utils.degree(edge_index[1], __num_nodes)[edge_index[1]], + ] + ) _temp_tensor: torch.Tensor = torch.pow(_temp_tensor, -0.5) _temp_tensor[torch.isinf(_temp_tensor)] = 0 return _temp_tensor[0] * _temp_tensor[1] def __init__( - self, edge_index: torch.LongTensor, - target_nodes_indexes: torch.LongTensor, - layer_wise_arguments: _typing.Sequence, - batch_size: _typing.Optional[int] = 1, num_workers: int = 0, - shuffle: bool = True, **kwargs + self, + edge_index: torch.LongTensor, + target_nodes_indexes: torch.LongTensor, + layer_wise_arguments: _typing.Sequence, + batch_size: _typing.Optional[int] = 1, + num_workers: int = 0, + shuffle: bool = True, + **kwargs ): super(LayerDependentImportanceSampler, self).__init__( torch_geometric.utils.add_remaining_self_loops(edge_index)[0], - target_nodes_indexes, layer_wise_arguments, batch_size, num_workers, shuffle, **kwargs + target_nodes_indexes, + layer_wise_arguments, + batch_size, + num_workers, + shuffle, + **kwargs ) self.__edge_weight: torch.Tensor = self.__compute_edge_weight(self._edge_index) - self.__integral_normalized_l_matrix: sp.csr_matrix = sp.csr_matrix(( - self.__edge_weight.numpy(), - (self._edge_index[1].numpy(), self._edge_index[0].numpy()) - )) - self.__integral_edges_indexes_sparse_matrix: sp.csr_matrix = sp.csr_matrix(( - np.arange(self._edge_index.size(1)), - (self._edge_index[1].numpy(), self._edge_index[0].numpy()) - )) + self.__integral_normalized_l_matrix: sp.csr_matrix = sp.csr_matrix( + ( + self.__edge_weight.numpy(), + (self._edge_index[1].numpy(), self._edge_index[0].numpy()), + ) + ) + self.__integral_edges_indexes_sparse_matrix: sp.csr_matrix = sp.csr_matrix( + ( + np.arange(self._edge_index.size(1)), + (self._edge_index[1].numpy(), self._edge_index[0].numpy()), + ) + ) def __sample_edges( - self, __current_layer_target_nodes_indexes: np.ndarray, - __top_layer_target_nodes_indexes: np.ndarray, sampled_source_nodes_budget: int + self, + __current_layer_target_nodes_indexes: np.ndarray, + __top_layer_target_nodes_indexes: np.ndarray, + sampled_source_nodes_budget: int, ) -> _typing.Tuple[np.ndarray, np.ndarray, np.ndarray]: """ @@ -280,33 +359,48 @@ class LayerDependentImportanceSampler(target_dependant_sampler.BasicLayerWiseTar corresponding probabilities for sampled_source_nodes_indexes ) """ - partial_l_matrix: sp.csr_matrix = ( - self.__integral_normalized_l_matrix[__current_layer_target_nodes_indexes, :] - ) - p: np.ndarray = np.array(np.sum(partial_l_matrix.multiply(partial_l_matrix), axis=0))[0] + partial_l_matrix: sp.csr_matrix = self.__integral_normalized_l_matrix[ + __current_layer_target_nodes_indexes, : + ] + p: np.ndarray = np.array( + np.sum(partial_l_matrix.multiply(partial_l_matrix), axis=0) + )[0] p: np.ndarray = p / np.sum(p) - _number_of_nodes_to_sample = np.min([np.sum(p > 0), sampled_source_nodes_budget]) - _selected_source_nodes: np.ndarray = np.unique(np.concatenate([ - np.random.choice( - p.size, _number_of_nodes_to_sample, replace=False, p=p - ), - __top_layer_target_nodes_indexes - ])) + _number_of_nodes_to_sample = np.min( + [np.sum(p > 0), sampled_source_nodes_budget] + ) + _selected_source_nodes: np.ndarray = np.unique( + np.concatenate( + [ + np.random.choice( + p.size, _number_of_nodes_to_sample, replace=False, p=p + ), + __top_layer_target_nodes_indexes, + ] + ) + ) _sampled_edges_indexes_sparse_matrix: sp.csr_matrix = ( - self.__integral_edges_indexes_sparse_matrix[__current_layer_target_nodes_indexes, :] + self.__integral_edges_indexes_sparse_matrix[ + __current_layer_target_nodes_indexes, : + ] ) _sampled_edges_indexes_sparse_matrix: sp.csc_matrix = ( _sampled_edges_indexes_sparse_matrix.tocsc()[:, _selected_source_nodes] ) - _sampled_edges_indexes: np.ndarray = np.unique(_sampled_edges_indexes_sparse_matrix.data) + _sampled_edges_indexes: np.ndarray = np.unique( + _sampled_edges_indexes_sparse_matrix.data + ) return _sampled_edges_indexes, _selected_source_nodes, p[_selected_source_nodes] def _sample_edges_for_layer( - self, __current_layer_target_nodes_indexes: torch.LongTensor, - __top_layer_target_nodes_indexes: torch.LongTensor, - layer_argument: _typing.Any, *args, **kwargs + self, + __current_layer_target_nodes_indexes: torch.LongTensor, + __top_layer_target_nodes_indexes: torch.LongTensor, + layer_argument: _typing.Any, + *args, + **kwargs ) -> _typing.Tuple[torch.LongTensor, _typing.Optional[torch.Tensor]]: """ Sample edges for one specific layer, expected to be implemented in subclass. @@ -331,34 +425,55 @@ class LayerDependentImportanceSampler(target_dependant_sampler.BasicLayerWiseTar edge_weight: the optional `edge_weight` for aggregation """ - __wrapped_result: _typing.Tuple[np.ndarray, np.ndarray, np.ndarray] = self.__sample_edges( + __wrapped_result: _typing.Tuple[ + np.ndarray, np.ndarray, np.ndarray + ] = self.__sample_edges( __current_layer_target_nodes_indexes.numpy(), __top_layer_target_nodes_indexes.numpy(), - layer_argument + layer_argument, ) _sampled_edges_indexes: torch.Tensor = torch.from_numpy(__wrapped_result[0]) _selected_source_nodes: torch.Tensor = torch.from_numpy(__wrapped_result[1]) - _selected_source_nodes_probabilities: torch.Tensor = torch.from_numpy(__wrapped_result[2]) + _selected_source_nodes_probabilities: torch.Tensor = torch.from_numpy( + __wrapped_result[2] + ) """ Multiply corresponding discount weights """ __selected_source_node_probability_mapping: _typing.Dict[int, float] = dict( - zip(_selected_source_nodes.tolist(), _selected_source_nodes_probabilities.tolist()) + zip( + _selected_source_nodes.tolist(), + _selected_source_nodes_probabilities.tolist(), + ) + ) + _selected_edges_weight: torch.Tensor = self.__edge_weight[ + _sampled_edges_indexes + ] + _selected_edges_weight: torch.Tensor = _selected_edges_weight / torch.tensor( + [ + __selected_source_node_probability_mapping.get( + _current_source_node_index + ) + for _current_source_node_index in self._edge_index[ + 0, _sampled_edges_indexes + ].tolist() + ] ) - _selected_edges_weight: torch.Tensor = self.__edge_weight[_sampled_edges_indexes] - _selected_edges_weight: torch.Tensor = _selected_edges_weight / torch.tensor([ - __selected_source_node_probability_mapping.get(_current_source_node_index) - for _current_source_node_index in self._edge_index[0, _sampled_edges_indexes].tolist() - ]) """ Normalize edge weight for selected edges by target nodes """ - for _current_target_node_index in self._edge_index[1, _sampled_edges_indexes].unique().tolist(): + for _current_target_node_index in ( + self._edge_index[1, _sampled_edges_indexes].unique().tolist() + ): _current_mask_for_selected_edges: torch.BoolTensor = ( - self._edge_index[1, _sampled_edges_indexes] == _current_target_node_index + self._edge_index[1, _sampled_edges_indexes] + == _current_target_node_index ) - _selected_edges_weight[_current_mask_for_selected_edges] = ( - _selected_edges_weight[_current_mask_for_selected_edges] / - torch.sum(_selected_edges_weight[_current_mask_for_selected_edges]) + _selected_edges_weight[ + _current_mask_for_selected_edges + ] = _selected_edges_weight[_current_mask_for_selected_edges] / torch.sum( + _selected_edges_weight[_current_mask_for_selected_edges] ) - _sampled_edges_indexes: _typing.Union[torch.LongTensor, torch.Tensor] = _sampled_edges_indexes + _sampled_edges_indexes: _typing.Union[ + torch.LongTensor, torch.Tensor + ] = _sampled_edges_indexes return _sampled_edges_indexes, _selected_edges_weight diff --git a/autogl/module/train/sampling/sampler/neighbor_sampler.py b/autogl/module/train/sampling/sampler/neighbor_sampler.py index 9d98b88..a9a7337 100644 --- a/autogl/module/train/sampling/sampler/neighbor_sampler.py +++ b/autogl/module/train/sampling/sampler/neighbor_sampler.py @@ -33,6 +33,7 @@ class NeighborSampler(TargetDependantSampler, _typing.Iterable): shuffle: whether to shuffle target nodes for mini-batches. """ + class _SequenceDataset(torch.utils.data.Dataset): def __init__(self, sequence): self.__sequence = sequence @@ -60,11 +61,14 @@ class NeighborSampler(TargetDependantSampler, _typing.Iterable): return temp_tensor[0] * temp_tensor[1] def __init__( - self, edge_index: torch.LongTensor, - target_nodes_indexes: torch.LongTensor, - sampling_sizes: _typing.Sequence[int], - batch_size: int = 1, num_workers: int = 0, - shuffle: bool = True, **kwargs + self, + edge_index: torch.LongTensor, + target_nodes_indexes: torch.LongTensor, + sampling_sizes: _typing.Sequence[int], + batch_size: int = 1, + num_workers: int = 0, + shuffle: bool = True, + **kwargs ): def is_deterministic(__cached: bool = bool(kwargs.get("cached", True))) -> bool: if not __cached: @@ -72,17 +76,25 @@ class NeighborSampler(TargetDependantSampler, _typing.Iterable): _deterministic: bool = True for _sampling_size in sampling_sizes: if type(_sampling_size) != int: - raise TypeError("The sampling_sizes argument must be a sequence of integer") + raise TypeError( + "The sampling_sizes argument must be a sequence of integer" + ) if _sampling_size >= 0: _deterministic = False break return _deterministic + self.__edge_weight: torch.Tensor = self.__compute_edge_weight(edge_index) self.__pyg_neighbor_sampler: torch_geometric.data.NeighborSampler = ( torch_geometric.data.NeighborSampler( - edge_index, list(sampling_sizes[::-1]), target_nodes_indexes, - transform=self._transform, batch_size=batch_size, - num_workers=num_workers, shuffle=shuffle, **kwargs + edge_index, + list(sampling_sizes[::-1]), + target_nodes_indexes, + transform=self._transform, + batch_size=batch_size, + num_workers=num_workers, + shuffle=shuffle, + **kwargs ) ) @@ -97,54 +109,78 @@ class NeighborSampler(TargetDependantSampler, _typing.Iterable): ] = None def _transform( - self, batch_size: int, n_id: torch.LongTensor, + self, + batch_size: int, + n_id: torch.LongTensor, adj_or_adj_list: _typing.Union[ _typing.Sequence[ - _typing.Tuple[torch.LongTensor, torch.LongTensor, _typing.Tuple[int, int]] + _typing.Tuple[ + torch.LongTensor, torch.LongTensor, _typing.Tuple[int, int] + ] ], - _typing.Tuple[torch.LongTensor, torch.LongTensor, _typing.Tuple[int, int]] - ] + _typing.Tuple[torch.LongTensor, torch.LongTensor, _typing.Tuple[int, int]], + ], ) -> TargetDependantSampledData: if ( - isinstance(adj_or_adj_list[0], _typing.Tuple) and - isinstance(adj_or_adj_list, _typing.Sequence) and - not isinstance(adj_or_adj_list, _typing.Tuple) + isinstance(adj_or_adj_list[0], _typing.Tuple) + and isinstance(adj_or_adj_list, _typing.Sequence) + and not isinstance(adj_or_adj_list, _typing.Tuple) ): return TargetDependantSampledData( [ - (current_layer[0], current_layer[1], self.__edge_weight[current_layer[1]]) + ( + current_layer[0], + current_layer[1], + self.__edge_weight[current_layer[1]], + ) for current_layer in adj_or_adj_list ], - (torch.arange(batch_size, dtype=torch.long).long(), n_id[:batch_size]), n_id + (torch.arange(batch_size, dtype=torch.long).long(), n_id[:batch_size]), + n_id, ) - elif isinstance(adj_or_adj_list, _typing.Tuple) and type(adj_or_adj_list[0]) == torch.Tensor: + elif ( + isinstance(adj_or_adj_list, _typing.Tuple) + and type(adj_or_adj_list[0]) == torch.Tensor + ): adj_or_adj_list: _typing.Tuple[ torch.LongTensor, torch.LongTensor, _typing.Tuple[int, int] ] = adj_or_adj_list return TargetDependantSampledData( - [(adj_or_adj_list[0], adj_or_adj_list[1], self.__edge_weight[adj_or_adj_list[1]])], - (torch.arange(batch_size, dtype=torch.long).long(), n_id[:batch_size]), n_id + [ + ( + adj_or_adj_list[0], + adj_or_adj_list[1], + self.__edge_weight[adj_or_adj_list[1]], + ) + ], + (torch.arange(batch_size, dtype=torch.long).long(), n_id[:batch_size]), + n_id, ) def __iter__(self): - if ( - self.__cached_sampled_data_list is not None and - isinstance(self.__cached_sampled_data_list, _typing.Sequence) + if self.__cached_sampled_data_list is not None and isinstance( + self.__cached_sampled_data_list, _typing.Sequence ): - return iter(torch.utils.data.DataLoader( - self._SequenceDataset(self.__cached_sampled_data_list), - collate_fn=lambda x: x[0] - )) + return iter( + torch.utils.data.DataLoader( + self._SequenceDataset(self.__cached_sampled_data_list), + collate_fn=lambda x: x[0], + ) + ) else: return iter(self.__pyg_neighbor_sampler) @classmethod def create_basic_sampler( - cls, edge_index: torch.LongTensor, - target_nodes_indexes: torch.LongTensor, - layer_wise_arguments: _typing.Sequence, - batch_size: int = 1, num_workers: int = 1, - shuffle: bool = True, *args, **kwargs + cls, + edge_index: torch.LongTensor, + target_nodes_indexes: torch.LongTensor, + layer_wise_arguments: _typing.Sequence, + batch_size: int = 1, + num_workers: int = 1, + shuffle: bool = True, + *args, + **kwargs ) -> TargetDependantSampler: """ A static factory method to create instance of :class:`NeighborSampler` @@ -172,6 +208,11 @@ class NeighborSampler(TargetDependantSampler, _typing.Iterable): whether to shuffle target nodes for mini-batches. """ return cls( - edge_index, target_nodes_indexes, layer_wise_arguments, - batch_size, num_workers, shuffle, **kwargs + edge_index, + target_nodes_indexes, + layer_wise_arguments, + batch_size, + num_workers, + shuffle, + **kwargs ) diff --git a/autogl/module/train/sampling/sampler/target_dependant_sampler.py b/autogl/module/train/sampling/sampler/target_dependant_sampler.py index f6b3547..679d544 100644 --- a/autogl/module/train/sampling/sampler/target_dependant_sampler.py +++ b/autogl/module/train/sampling/sampler/target_dependant_sampler.py @@ -38,18 +38,18 @@ class TargetDependantSampledData: The stored sequence of tuple `( edge_index_for_sampled_graph, edge_id_in_integral_graph, (optional)edge_weight )`. """ + class _LayerSampledEdgeData: def __init__( - self, edge_index_for_sampled_graph: torch.Tensor, - edge_id_in_integral_graph: torch.Tensor, - edge_weight: _typing.Optional[torch.Tensor] + self, + edge_index_for_sampled_graph: torch.Tensor, + edge_id_in_integral_graph: torch.Tensor, + edge_weight: _typing.Optional[torch.Tensor], ): self.__edge_index_for_sampled_graph: torch.Tensor = ( edge_index_for_sampled_graph ) - self.__edge_id_in_integral_graph: torch.Tensor = ( - edge_id_in_integral_graph - ) + self.__edge_id_in_integral_graph: torch.Tensor = edge_id_in_integral_graph self.__edge_weight: _typing.Optional[torch.Tensor] = edge_weight @property @@ -61,9 +61,7 @@ class TargetDependantSampledData: @property def edge_id_in_integral_graph(self) -> torch.LongTensor: - edge_id_in_integral_graph: _typing.Any = ( - self.__edge_id_in_integral_graph - ) + edge_id_in_integral_graph: _typing.Any = self.__edge_id_in_integral_graph return edge_id_in_integral_graph @property @@ -82,9 +80,9 @@ class TargetDependantSampledData: return indexes_in_integral_graph def __init__( - self, - indexes_in_sampled_graph: torch.Tensor, - indexes_in_integral_graph: torch.Tensor, + self, + indexes_in_sampled_graph: torch.Tensor, + indexes_in_integral_graph: torch.Tensor, ): self.__indexes_in_sampled_graph: torch.Tensor = indexes_in_sampled_graph self.__indexes_in_integral_graph: torch.Tensor = indexes_in_integral_graph @@ -105,12 +103,12 @@ class TargetDependantSampledData: return self.__sampled_edges_for_layers def __init__( - self, - sampled_edges_for_layers: _typing.Sequence[ - _typing.Tuple[torch.Tensor, torch.Tensor, _typing.Optional[torch.Tensor]] - ], - target_nodes_indexes: _typing.Tuple[torch.Tensor, torch.Tensor], - all_sampled_nodes_indexes: torch.Tensor + self, + sampled_edges_for_layers: _typing.Sequence[ + _typing.Tuple[torch.Tensor, torch.Tensor, _typing.Optional[torch.Tensor]] + ], + target_nodes_indexes: _typing.Tuple[torch.Tensor, torch.Tensor], + all_sampled_nodes_indexes: torch.Tensor, ): self.__sampled_edges_for_layers: _typing.Sequence[ TargetDependantSampledData._LayerSampledEdgeData @@ -128,13 +126,18 @@ class TargetDependantSampler(torch.utils.data.DataLoader, _typing.Iterable): """ An abstract base class for various target-dependent sampler """ + @classmethod def create_basic_sampler( - cls, edge_index: torch.LongTensor, - target_nodes_indexes: torch.LongTensor, - layer_wise_arguments: _typing.Sequence, - batch_size: int = 1, num_workers: int = 0, - shuffle: bool = True, *args, **kwargs + cls, + edge_index: torch.LongTensor, + target_nodes_indexes: torch.LongTensor, + layer_wise_arguments: _typing.Sequence, + batch_size: int = 1, + num_workers: int = 0, + shuffle: bool = True, + *args, + **kwargs ) -> "TargetDependantSampler": """ :param edge_index: edge index of integral graph @@ -148,7 +151,7 @@ class TargetDependantSampler(torch.utils.data.DataLoader, _typing.Iterable): :return: instance of TargetDependantSampler """ raise NotImplementedError - + def __iter__(self): return super(TargetDependantSampler, self).__iter__() @@ -175,12 +178,16 @@ class BasicLayerWiseTargetDependantSampler(TargetDependantSampler): kwargs: remaining keyword arguments """ + def __init__( - self, edge_index: torch.LongTensor, - target_nodes_indexes: torch.LongTensor, - layer_wise_arguments: _typing.Sequence, - batch_size: _typing.Optional[int] = 1, num_workers: int = 0, - shuffle: bool = True, **kwargs + self, + edge_index: torch.LongTensor, + target_nodes_indexes: torch.LongTensor, + layer_wise_arguments: _typing.Sequence, + batch_size: _typing.Optional[int] = 1, + num_workers: int = 0, + shuffle: bool = True, + **kwargs ): self._edge_index: torch.LongTensor = edge_index self.__layer_wise_arguments: _typing.Sequence = layer_wise_arguments @@ -188,17 +195,24 @@ class BasicLayerWiseTargetDependantSampler(TargetDependantSampler): del kwargs["collate_fn"] super(BasicLayerWiseTargetDependantSampler, self).__init__( target_nodes_indexes.unique().numpy(), - batch_size, shuffle, num_workers=num_workers, - collate_fn=self._collate_fn, **kwargs + batch_size, + shuffle, + num_workers=num_workers, + collate_fn=self._collate_fn, + **kwargs ) @classmethod def create_basic_sampler( - cls, edge_index: torch.LongTensor, - target_nodes_indexes: torch.LongTensor, - layer_wise_arguments: _typing.Sequence, - batch_size: int = 1, num_workers: int = 0, - shuffle: bool = True, *args, **kwargs + cls, + edge_index: torch.LongTensor, + target_nodes_indexes: torch.LongTensor, + layer_wise_arguments: _typing.Sequence, + batch_size: int = 1, + num_workers: int = 0, + shuffle: bool = True, + *args, + **kwargs ) -> TargetDependantSampler: """ :param edge_index: edge index of integral graph @@ -212,14 +226,22 @@ class BasicLayerWiseTargetDependantSampler(TargetDependantSampler): :return: instance of TargetDependantSampler """ return BasicLayerWiseTargetDependantSampler( - edge_index, target_nodes_indexes, layer_wise_arguments, - batch_size, num_workers, shuffle, **kwargs + edge_index, + target_nodes_indexes, + layer_wise_arguments, + batch_size, + num_workers, + shuffle, + **kwargs ) def _sample_edges_for_layer( - self, __current_layer_target_nodes_indexes: torch.LongTensor, - __top_layer_target_nodes_indexes: torch.LongTensor, - layer_argument: _typing.Any, *args, **kwargs + self, + __current_layer_target_nodes_indexes: torch.LongTensor, + __top_layer_target_nodes_indexes: torch.LongTensor, + layer_argument: _typing.Any, + *args, + **kwargs ) -> _typing.Tuple[torch.LongTensor, _typing.Optional[torch.Tensor]]: """ Sample edges for one specific layer, expected to be implemented in subclass. @@ -247,17 +269,21 @@ class BasicLayerWiseTargetDependantSampler(TargetDependantSampler): raise NotImplementedError def _collate_fn( - self, top_layer_target_nodes_indexes_list: _typing.List[int] + self, top_layer_target_nodes_indexes_list: _typing.List[int] ) -> TargetDependantSampledData: - return self.__sample_layers(torch.tensor(top_layer_target_nodes_indexes_list).unique()) + return self.__sample_layers( + torch.tensor(top_layer_target_nodes_indexes_list).unique() + ) def __sample_layers( - self, __top_layer_target_nodes_indexes: torch.LongTensor + self, __top_layer_target_nodes_indexes: torch.LongTensor ) -> TargetDependantSampledData: sampled_edges_for_layers: _typing.List[ _typing.Tuple[torch.LongTensor, _typing.Optional[torch.Tensor]] ] = list() - __current_layer_target_nodes_indexes: torch.LongTensor = __top_layer_target_nodes_indexes + __current_layer_target_nodes_indexes: torch.LongTensor = ( + __top_layer_target_nodes_indexes + ) " Reverse self.__layer_wise_arguments from bottom-up to top-down " for layer_argument in self.__layer_wise_arguments[::-1]: current_layer_result: _typing.Tuple[ @@ -265,11 +291,11 @@ class BasicLayerWiseTargetDependantSampler(TargetDependantSampler): ] = self._sample_edges_for_layer( __current_layer_target_nodes_indexes, __top_layer_target_nodes_indexes, - layer_argument - ) - __source_nodes_indexes_for_current_layer: torch.Tensor = ( - self._edge_index[0, current_layer_result[0]] + layer_argument, ) + __source_nodes_indexes_for_current_layer: torch.Tensor = self._edge_index[ + 0, current_layer_result[0] + ] __current_layer_target_nodes_indexes: torch.LongTensor = ( __source_nodes_indexes_for_current_layer.unique() ) @@ -285,49 +311,66 @@ class BasicLayerWiseTargetDependantSampler(TargetDependantSampler): for current_layer_result in sampled_edges_for_layers ] ).unique() - __sampled_nodes_in_sub_graph_mapping: _typing.Dict[int, int] = dict(list(zip( - sampled_nodes_in_sub_graph.tolist(), - range(sampled_nodes_in_sub_graph.size(0)) - ))) + __sampled_nodes_in_sub_graph_mapping: _typing.Dict[int, int] = dict( + list( + zip( + sampled_nodes_in_sub_graph.tolist(), + range(sampled_nodes_in_sub_graph.size(0)), + ) + ) + ) __sampled_edge_index_for_layers_in_sub_graph: _typing.Sequence[torch.Tensor] = [ - torch.stack([ - torch.tensor( - [ - __sampled_nodes_in_sub_graph_mapping.get(node_index) - for node_index in self._edge_index[0, current_layer_result[0]].tolist() - ] - ), - torch.tensor( - [ - __sampled_nodes_in_sub_graph_mapping.get(node_index) - for node_index in self._edge_index[1, current_layer_result[0]].tolist() - ] - ), - ]) + torch.stack( + [ + torch.tensor( + [ + __sampled_nodes_in_sub_graph_mapping.get(node_index) + for node_index in self._edge_index[ + 0, current_layer_result[0] + ].tolist() + ] + ), + torch.tensor( + [ + __sampled_nodes_in_sub_graph_mapping.get(node_index) + for node_index in self._edge_index[ + 1, current_layer_result[0] + ].tolist() + ] + ), + ] + ) for current_layer_result in sampled_edges_for_layers ] return TargetDependantSampledData( [ - (temp_tuple[0], temp_tuple[1][0], temp_tuple[1][1]) for temp_tuple - in zip(__sampled_edge_index_for_layers_in_sub_graph, sampled_edges_for_layers) + (temp_tuple[0], temp_tuple[1][0], temp_tuple[1][1]) + for temp_tuple in zip( + __sampled_edge_index_for_layers_in_sub_graph, + sampled_edges_for_layers, + ) ], ( torch.tensor( [ - __sampled_nodes_in_sub_graph_mapping.get(current_target_node_index_in_integral_data) + __sampled_nodes_in_sub_graph_mapping.get( + current_target_node_index_in_integral_data + ) for current_target_node_index_in_integral_data in __top_layer_target_nodes_indexes.tolist() - if current_target_node_index_in_integral_data in __sampled_nodes_in_sub_graph_mapping + if current_target_node_index_in_integral_data + in __sampled_nodes_in_sub_graph_mapping ] ).long(), # Remap torch.tensor( [ current_target_node_index_in_integral_data for current_target_node_index_in_integral_data in __top_layer_target_nodes_indexes.tolist() - if current_target_node_index_in_integral_data in __sampled_nodes_in_sub_graph_mapping + if current_target_node_index_in_integral_data + in __sampled_nodes_in_sub_graph_mapping ] - ).long() + ).long(), ), - sampled_nodes_in_sub_graph + sampled_nodes_in_sub_graph, ) diff --git a/autogl/solver/base.py b/autogl/solver/base.py index 75083d9..f25f4bc 100644 --- a/autogl/solver/base.py +++ b/autogl/solver/base.py @@ -298,23 +298,65 @@ class BaseSolver: if nas_algorithms is None and nas_estimators is None and nas_spaces is None: self.nas_algorithms = self.nas_estimators = self.nas_spaces = None return - assert None not in [nas_algorithms, nas_estimators, nas_spaces], "The algorithms, estimators and spaces should all be set" - - nas_algorithms = nas_algorithms if isinstance(nas_algorithms, (list, tuple)) else [nas_algorithms] - nas_spaces = nas_spaces if isinstance(nas_spaces, (list, tuple)) else [nas_spaces] - nas_estimators = nas_estimators if isinstance(nas_estimators, (list, tuple)) else [nas_estimators] + assert None not in [ + nas_algorithms, + nas_estimators, + nas_spaces, + ], "The algorithms, estimators and spaces should all be set" + + nas_algorithms = ( + nas_algorithms + if isinstance(nas_algorithms, (list, tuple)) + else [nas_algorithms] + ) + nas_spaces = ( + nas_spaces if isinstance(nas_spaces, (list, tuple)) else [nas_spaces] + ) + nas_estimators = ( + nas_estimators + if isinstance(nas_estimators, (list, tuple)) + else [nas_estimators] + ) # parse all str elements - nas_algorithms = [algo if not isinstance(algo, str) else NAS_ALGO_DICT[algo]() for algo in nas_algorithms] - nas_spaces = [space if not isinstance(space, str) else NAS_SPACE_DICT[space]() for space in nas_spaces] - nas_estimators = [estimator if not isinstance(estimator, str) else NAS_ESTIMATOR_DICT[estimator]() for estimator in nas_estimators] - - max_number = max([len(x) for x in [nas_algorithms, nas_spaces, nas_estimators]]) - assert all([len(x) in [1, max_number] for x in [nas_algorithms, nas_spaces, nas_estimators]]), "lengths of algorithms/spaces/estimators do not match!" + nas_algorithms = [ + algo if not isinstance(algo, str) else NAS_ALGO_DICT[algo]() + for algo in nas_algorithms + ] + nas_spaces = [ + space if not isinstance(space, str) else NAS_SPACE_DICT[space]() + for space in nas_spaces + ] + nas_estimators = [ + estimator + if not isinstance(estimator, str) + else NAS_ESTIMATOR_DICT[estimator]() + for estimator in nas_estimators + ] - self.nas_algorithms = [deepcopy(nas_algorithms) for _ in range(max_number)] if len(nas_algorithms) == 1 and max_number > 1 else nas_algorithms - self.nas_spaces = [deepcopy(nas_spaces) for _ in range(max_number)] if len(nas_spaces) == 1 and max_number > 1 else nas_spaces - self.nas_estimators = [deepcopy(nas_estimators) for _ in range(max_number)] if len(nas_estimators) == 1 and max_number > 1 else nas_estimators + max_number = max([len(x) for x in [nas_algorithms, nas_spaces, nas_estimators]]) + assert all( + [ + len(x) in [1, max_number] + for x in [nas_algorithms, nas_spaces, nas_estimators] + ] + ), "lengths of algorithms/spaces/estimators do not match!" + + self.nas_algorithms = ( + [deepcopy(nas_algorithms) for _ in range(max_number)] + if len(nas_algorithms) == 1 and max_number > 1 + else nas_algorithms + ) + self.nas_spaces = ( + [deepcopy(nas_spaces) for _ in range(max_number)] + if len(nas_spaces) == 1 and max_number > 1 + else nas_spaces + ) + self.nas_estimators = ( + [deepcopy(nas_estimators) for _ in range(max_number)] + if len(nas_estimators) == 1 and max_number > 1 + else nas_estimators + ) return self diff --git a/autogl/solver/classifier/graph_classifier.py b/autogl/solver/classifier/graph_classifier.py index e77d00d..f7751a5 100644 --- a/autogl/solver/classifier/graph_classifier.py +++ b/autogl/solver/classifier/graph_classifier.py @@ -90,9 +90,9 @@ class AutoGraphClassifier(BaseClassifier): super().__init__( feature_module=feature_module, graph_models=graph_models, - nas_algorithms=None, #nas_algorithms, - nas_spaces=None, #nas_spaces, - nas_estimators=None, #nas_estimators, + nas_algorithms=None, # nas_algorithms, + nas_spaces=None, # nas_spaces, + nas_estimators=None, # nas_estimators, hpo_module=hpo_module, ensemble_module=ensemble_module, max_evals=max_evals, diff --git a/autogl/solver/classifier/node_classifier.py b/autogl/solver/classifier/node_classifier.py index da676bc..cc8c2c1 100644 --- a/autogl/solver/classifier/node_classifier.py +++ b/autogl/solver/classifier/node_classifier.py @@ -43,13 +43,13 @@ class AutoNodeClassifier(BaseClassifier): graph_models: list of autogl.module.model.BaseModel or list of str The (name of) models to be optimized as backbone. Default ``['gat', 'gcn']``. - + nas_algorithms: (list of) autogl.module.nas.algorithm.BaseNAS or str (Optional) The (name of) nas algorithms used. Default ``None``. - + nas_spaces: (list of) autogl.module.nas.space.BaseSpace or str (Optional) The (name of) nas spaces used. Default ``None``. - + nas_estimators: (list of) autogl.module.nas.estimator.BaseEstimator or str (Optional) The (name of) nas estimators used. Default ``None``. @@ -167,7 +167,6 @@ class AutoNodeClassifier(BaseClassifier): loss=loss, feval=feval, device=device, - ) self.graph_model_list.append(model) else: @@ -365,7 +364,11 @@ class AutoNodeClassifier(BaseClassifier): loss="nll_loss" if not hasattr(dataset, "loss") else dataset.loss, ) - assert not isinstance(self._default_trainer, list) or len(self.nas_algorithms) == len(self._default_trainer) - len(self.graph_model_list), "length of default trainer should match total graph models and nas models passed" + assert not isinstance(self._default_trainer, list) or len( + self.nas_algorithms + ) == len(self._default_trainer) - len( + self.graph_model_list + ), "length of default trainer should match total graph models and nas models passed" # perform nas and add them to model list idx_trainer = len(self.graph_model_list) @@ -384,7 +387,9 @@ class AutoNodeClassifier(BaseClassifier): model=model, num_features=self.dataset[0].x.shape[1], num_classes=self.dataset.num_classes, - loss="nll_loss" if not hasattr(dataset, "loss") else dataset.loss, + loss="nll_loss" + if not hasattr(dataset, "loss") + else dataset.loss, feval=evaluator_list, device=self.runtime_device, init=False, @@ -395,10 +400,11 @@ class AutoNodeClassifier(BaseClassifier): trainer.update_parameters( num_classes=self.dataset.num_classes, num_features=self.dataset[0].x.shape[1], - loss="nll_loss" if not hasattr(dataset, "loss") else dataset.loss, + loss="nll_loss" + if not hasattr(dataset, "loss") + else dataset.loss, feval=evaluator_list, device=self.runtime_device, - ) self.graph_model_list.append(trainer) @@ -814,25 +820,29 @@ class AutoNodeClassifier(BaseClassifier): if ensemble_dict is not None: name = ensemble_dict.pop("name") solver.set_ensemble_module(name, **ensemble_dict) - + nas_dict = path_or_dict.pop("nas", None) if nas_dict is not None: keys: set = set(nas_dict.keys()) - needed = {'space', 'algorithm', 'estimator'} + needed = {"space", "algorithm", "estimator"} if keys != needed: - LOGGER.error('Key mismatch, we need %s, you give %s', needed, keys) - raise KeyError('Key mismatch, we need %s, you give %s' % (needed, keys)) + LOGGER.error("Key mismatch, we need %s, you give %s", needed, keys) + raise KeyError("Key mismatch, we need %s, you give %s" % (needed, keys)) spaces, algorithms, estimators = [], [], [] - for container, indexer, k in zip([spaces, algorithms, estimators], [NAS_SPACE_DICT, NAS_ALGO_DICT, NAS_ESTIMATOR_DICT], ['space', 'algorithm', 'estimator']): + for container, indexer, k in zip( + [spaces, algorithms, estimators], + [NAS_SPACE_DICT, NAS_ALGO_DICT, NAS_ESTIMATOR_DICT], + ["space", "algorithm", "estimator"], + ): configs = nas_dict[k] if isinstance(configs, list): for item in configs: - container.append(indexer[item.pop('name')](**item)) + container.append(indexer[item.pop("name")](**item)) else: - container.append(indexer[configs.pop('name')](**configs)) - + container.append(indexer[configs.pop("name")](**configs)) + solver.set_nas_module(algorithms, spaces, estimators) return solver diff --git a/autogl/solver/utils.py b/autogl/solver/utils.py index 03fc36b..86843e5 100644 --- a/autogl/solver/utils.py +++ b/autogl/solver/utils.py @@ -52,7 +52,9 @@ class LeaderBoard: if field in self.keys and not field == "name": self.major_field = field else: - LOGGER.warning(f"Field [{field}] NOT found in the current LeaderBoard, will ignore.") + LOGGER.warning( + f"Field [{field}] NOT found in the current LeaderBoard, will ignore." + ) def insert_model_performance(self, name, performance) -> None: """ @@ -144,10 +146,10 @@ class LeaderBoard: """ top_k: int = top_k if top_k > 0 else len(self.perform_dict) - ''' + """ reindex self.__performance_data_frame to ensure the columns of name and representation are in left-side of the data frame - ''' + """ _columns = self.perform_dict.columns.tolist() maxcolwidths: _typing.List[_typing.Optional[int]] = [] if "name" in _columns: @@ -157,18 +159,19 @@ class LeaderBoard: self.perform_dict = self.perform_dict[_columns] sorted_performance_df: pd.DataFrame = self.perform_dict.sort_values( - self.major_field, - ascending=not self.is_higher_better[self.major_field] + self.major_field, ascending=not self.is_higher_better[self.major_field] ) sorted_performance_df = sorted_performance_df.head(top_k) from tabulate import tabulate + _columns = sorted_performance_df.columns.tolist() maxcolwidths.extend([None for _ in range(len(_columns) - len(maxcolwidths))]) print( tabulate( list(zip(*[sorted_performance_df[column] for column in _columns])), - headers=_columns, tablefmt="grid" + headers=_columns, + tablefmt="grid", ) ) diff --git a/autogl/utils/device.py b/autogl/utils/device.py index b331ace..05299e4 100644 --- a/autogl/utils/device.py +++ b/autogl/utils/device.py @@ -1,24 +1,27 @@ import torch from typing import Union + def get_device(device: Union[str, torch.device]): """ Get device of passed argument. Will return a torch.device based on passed arguments. Can parse auto, cpu, gpu, cpu:x, gpu:x, etc. If auto is given, will automatically find available devices. - + Parameters ---------- device: ``str`` or ``torch.device`` The device to parse. If ``auto`` if given, will determine automatically. - + Returns ------- device: ``torch.device`` The parsed device. """ - assert isinstance(device, (str, torch.device)), "Only support device of str or torch.device, get {} instead".format(device) - if device == 'auto': - device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + assert isinstance( + device, (str, torch.device) + ), "Only support device of str or torch.device, get {} instead".format(device) + if device == "auto": + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") return torch.device(device)