From 1edc5692ed083e050eab28b67665672ca72fa4fb Mon Sep 17 00:00:00 2001 From: null Date: Thu, 10 Jun 2021 02:44:00 +0800 Subject: [PATCH] Modify the GCN model, Improve sampling speed, GC before training Modify the GCN model to support diverse dropout probabilities in different GCN layers. Improve the sampling speed for BasicTargetDependentSampler base class. Apply GC before training for every instance of trainer. --- autogl/module/model/gcn.py | 43 ++++++++++++++++--- .../node_classification_sampled_trainer.py | 6 +++ .../sampler/target_dependant_sampler.py | 2 +- 3 files changed, 45 insertions(+), 6 deletions(-) diff --git a/autogl/module/model/gcn.py b/autogl/module/model/gcn.py index 1f970ac..f3abd15 100644 --- a/autogl/module/model/gcn.py +++ b/autogl/module/model/gcn.py @@ -72,33 +72,66 @@ class GCN(SequentialGraphNeuralNetwork): num_features: int, num_classes: int, hidden_features: _typing.Sequence[int], - dropout: float, + dropout: _typing.Union[float, _typing.Sequence[_typing.Optional[float]]], activation_name: str, add_self_loops: bool = True, normalize: bool = True ): + if isinstance(dropout, _typing.Sequence): + if len(dropout) != len(hidden_features) + 1: + raise TypeError( + "When the dropout argument is a sequence, " + "The sequence length must equal to the number of layers to construct." + ) + for _dropout in dropout: + if _dropout is not None and type(_dropout) != float: + raise TypeError( + "When the dropout argument is a sequence, " + "every item in the sequence must be float or None" + ) + dropout_list: _typing.Sequence[_typing.Optional[float]] = dropout + elif type(dropout) == float: + if dropout < 0: + dropout = 0 + if dropout > 1: + dropout = 1 + dropout_list: _typing.Sequence[_typing.Optional[float]] = [ + dropout for _ in range(len(hidden_features) + 1) + ] + else: + raise TypeError( + "The provided dropout argument must be a float " + "or a sequence in which each item is either float or None." + ) super().__init__() if len(hidden_features) == 0: self.__sequential_module_list: torch.nn.ModuleList = torch.nn.ModuleList( - (self._GCNLayer(num_features, num_classes, add_self_loops, normalize),) + ( + self._GCNLayer( + num_features, num_classes, add_self_loops, normalize, + dropout_probability=dropout_list[0] + ), + ) ) else: self.__sequential_module_list: torch.nn.ModuleList = torch.nn.ModuleList() self.__sequential_module_list.append(self._GCNLayer( num_features, hidden_features[0], add_self_loops, - normalize, activation_name, dropout + 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_module_list.append(self._GCNLayer( hidden_features[hidden_feature_index], hidden_features[hidden_feature_index + 1], - add_self_loops, normalize, activation_name, dropout + add_self_loops, normalize, activation_name, + dropout_list[hidden_feature_index + 1] )) else: self.__sequential_module_list.append(self._GCNLayer( hidden_features[hidden_feature_index], num_classes, - add_self_loops, normalize + add_self_loops, normalize, + dropout_list[-1] )) def decode(self, x: torch.Tensor) -> torch.Tensor: diff --git a/autogl/module/train/node_classification_trainer/node_classification_sampled_trainer.py b/autogl/module/train/node_classification_trainer/node_classification_sampled_trainer.py index 1f5323f..1e611c7 100644 --- a/autogl/module/train/node_classification_trainer/node_classification_sampled_trainer.py +++ b/autogl/module/train/node_classification_trainer/node_classification_sampled_trainer.py @@ -319,6 +319,8 @@ class NodeClassificationGraphSAINTTrainer(BaseNodeClassificationTrainer): :param dataset: :param keep_valid_result: Whether to save the validation result after training """ + import gc + gc.collect() data = dataset[0] self.__train_only(data) if keep_valid_result: @@ -879,6 +881,8 @@ class NodeClassificationLayerDependentImportanceSamplingTrainer(BaseNodeClassifi :param dataset: :param keep_valid_result: Whether to save the validation result after training """ + import gc + gc.collect() data = dataset[0] self.__train_only(data) if keep_valid_result: @@ -1387,6 +1391,8 @@ class NodeClassificationNeighborSamplingTrainer(BaseNodeClassificationTrainer): :param dataset: :param keep_valid_result: Whether to save the validation result after training """ + import gc + gc.collect() data = dataset[0] self.__train_only(data) if keep_valid_result: diff --git a/autogl/module/train/sampling/sampler/target_dependant_sampler.py b/autogl/module/train/sampling/sampler/target_dependant_sampler.py index db51654..a506323 100644 --- a/autogl/module/train/sampling/sampler/target_dependant_sampler.py +++ b/autogl/module/train/sampling/sampler/target_dependant_sampler.py @@ -138,7 +138,7 @@ class BasicLayerWiseTargetDependantSampler(TargetDependantSampler): if "collate_fn" in kwargs: del kwargs["collate_fn"] super(BasicLayerWiseTargetDependantSampler, self).__init__( - target_nodes_indexes.unique().tolist(), + target_nodes_indexes.unique().numpy(), batch_size, shuffle, num_workers=num_workers, collate_fn=self._collate_fn, **kwargs )