black format

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",

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

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 *

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,

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
+        return data

autogl/module/feature/base.py

@@ -169,4 +169,4 @@ class TransformWrapper(BaseFeature):
             return self
 
     def _transform(self, data=None):
-        return self._func(data)
+        return self._func(data)

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",
 ]

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
+        return data

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
-    """
+    """

autogl/module/feature/graph/base.py

@@ -16,4 +16,4 @@ class BaseGraph(BaseFeature):
             data.gf = torch.FloatTensor([[]])
 
     def _postprocess(self, data):
-        pass
+        pass

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
+        return data

autogl/module/feature/graph/nx.py

@@ -180,4 +180,4 @@ class NxIsEulerian(NxGraph):
     pass
 
 
-# till algorithms.flows
+# till algorithms.flows

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
+        return data

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"])

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

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)

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)

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)

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

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"]

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

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

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

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
+    def _infer(self, mask="train"):
+        metric, loss = self.estimator.infer(self.arch._model, self.dataset, mask=mask)
+        return metric[0], loss

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"]

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
 

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

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

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",
+]

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.

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)
+        # return AutoGCN(self.input_dim, self.output_dim, device)
+        return self.wrap(device).fix(selection)

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)

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))

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)

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 = []

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")

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):
         """

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),
             }
         )
 

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
         )

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

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
         )

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,
         )

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
 

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,

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

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",
             )
         )
 

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)