From 7b551961e90f5042d9b91d92c083f3f09dd9dbdd Mon Sep 17 00:00:00 2001 From: Frozenmad Date: Tue, 23 Nov 2021 10:46:46 +0800 Subject: [PATCH 01/14] Update LICENSE update license with proper credit --- LICENSE | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/LICENSE b/LICENSE index 7a4a3ea..e29bff9 100644 --- a/LICENSE +++ b/LICENSE @@ -187,7 +187,7 @@ same "printed page" as the copyright notice for easier identification within third-party archives. - Copyright [yyyy] [name of copyright owner] + Copyright 2020-2021 AGLTeam THUMNLab and contributors Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -199,4 +199,4 @@ distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and - limitations under the License. \ No newline at end of file + limitations under the License. From bf7048e877c7ba692668b64840df367a0d2f1982 Mon Sep 17 00:00:00 2001 From: generall Date: Sun, 5 Dec 2021 11:20:37 +0800 Subject: [PATCH 02/14] add gasso --- autogl/module/nas/algorithm/__init__.py | 7 +- autogl/module/nas/algorithm/enas.py | 1 + autogl/module/nas/algorithm/gasso.py | 131 ++++++ autogl/module/nas/space/__init__.py | 6 + autogl/module/nas/space/gasso.py | 308 ++++++++++++++ .../module/nas/space/gasso_space/__init__.py | 23 ++ .../module/nas/space/gasso_space/arma_conv.py | 133 ++++++ .../module/nas/space/gasso_space/cheb_conv.py | 157 +++++++ .../module/nas/space/gasso_space/edge_conv.py | 123 ++++++ .../module/nas/space/gasso_space/gat_conv.py | 194 +++++++++ .../module/nas/space/gasso_space/gcn_conv.py | 200 +++++++++ .../module/nas/space/gasso_space/gin_conv.py | 157 +++++++ autogl/module/nas/space/gasso_space/inits.py | 56 +++ .../space/gasso_space/message_passing.jinja | 153 +++++++ .../nas/space/gasso_space/message_passing.py | 389 ++++++++++++++++++ .../module/nas/space/gasso_space/sage_conv.py | 92 +++++ .../nas/space/gasso_space/utils/__init__.py | 0 .../nas/space/gasso_space/utils/helpers.py | 7 + .../nas/space/gasso_space/utils/inspector.py | 86 ++++ .../module/nas/space/gasso_space/utils/jit.py | 19 + .../nas/space/gasso_space/utils/typing.py | 107 +++++ configs/nodeclf_nas_gasso.yml | 40 ++ examples/gasso_test.py | 24 ++ test/nas/node_classification.py | 16 +- 24 files changed, 2416 insertions(+), 13 deletions(-) create mode 100644 autogl/module/nas/algorithm/gasso.py create mode 100644 autogl/module/nas/space/gasso.py create mode 100644 autogl/module/nas/space/gasso_space/__init__.py create mode 100644 autogl/module/nas/space/gasso_space/arma_conv.py create mode 100644 autogl/module/nas/space/gasso_space/cheb_conv.py create mode 100644 autogl/module/nas/space/gasso_space/edge_conv.py create mode 100644 autogl/module/nas/space/gasso_space/gat_conv.py create mode 100644 autogl/module/nas/space/gasso_space/gcn_conv.py create mode 100644 autogl/module/nas/space/gasso_space/gin_conv.py create mode 100644 autogl/module/nas/space/gasso_space/inits.py create mode 100644 autogl/module/nas/space/gasso_space/message_passing.jinja create mode 100644 autogl/module/nas/space/gasso_space/message_passing.py create mode 100644 autogl/module/nas/space/gasso_space/sage_conv.py create mode 100644 autogl/module/nas/space/gasso_space/utils/__init__.py create mode 100644 autogl/module/nas/space/gasso_space/utils/helpers.py create mode 100644 autogl/module/nas/space/gasso_space/utils/inspector.py create mode 100644 autogl/module/nas/space/gasso_space/utils/jit.py create mode 100644 autogl/module/nas/space/gasso_space/utils/typing.py create mode 100644 configs/nodeclf_nas_gasso.yml create mode 100644 examples/gasso_test.py diff --git a/autogl/module/nas/algorithm/__init__.py b/autogl/module/nas/algorithm/__init__.py index 5056d17..50e9233 100644 --- a/autogl/module/nas/algorithm/__init__.py +++ b/autogl/module/nas/algorithm/__init__.py @@ -29,7 +29,9 @@ from .darts import Darts from .enas import Enas from .random_search import RandomSearch from .rl import RL, GraphNasRL - +from ..backend import * +if not is_dgl(): + from .gasso import Gasso def build_nas_algo_from_name(name: str) -> BaseNAS: """ @@ -53,3 +55,6 @@ def build_nas_algo_from_name(name: str) -> BaseNAS: __all__ = ["BaseNAS", "Darts", "Enas", "RandomSearch", "RL", "GraphNasRL"] + +if not is_dgl(): + __all__.append("Gasso") \ No newline at end of file diff --git a/autogl/module/nas/algorithm/enas.py b/autogl/module/nas/algorithm/enas.py index 56e0232..36fac6e 100644 --- a/autogl/module/nas/algorithm/enas.py +++ b/autogl/module/nas/algorithm/enas.py @@ -101,6 +101,7 @@ class Enas(BaseNAS): def search(self, space: BaseSpace, dset, estimator): self.model = space self.dataset = dset # .to(self.device) + print(dir(dset)) self.estimator = estimator # replace choice self.nas_modules = [] diff --git a/autogl/module/nas/algorithm/gasso.py b/autogl/module/nas/algorithm/gasso.py new file mode 100644 index 0000000..28ec33f --- /dev/null +++ b/autogl/module/nas/algorithm/gasso.py @@ -0,0 +1,131 @@ +# "Graph differentiable architecture search with structure optimization" NeurIPS 21' + +import logging + +import torch +import torch.optim +import torch.nn as nn +import torch.nn.functional as F + +from . import register_nas_algo +from .base import BaseNAS +from ..estimator.base import BaseEstimator +from ..space import BaseSpace +from ..utils import replace_layer_choice, replace_input_choice +from ...model.base import BaseModel + +from torch.autograd import Variable +import numpy as np +import time +import copy +import torch.optim as optim +import scipy.sparse as sp + +_logger = logging.getLogger(__name__) + +@register_nas_algo("gasso") +class Gasso(BaseNAS): + def __init__( + self, + num_epochs=250, + warmup_epochs=10, + model_lr=0.01, + model_wd=1e-4, + arch_lr = 0.03, + stru_lr = 0.04, + lamb = 0.6, + device="cuda", + ): + super().__init__(device=device) + self.device = device + self.num_epochs = num_epochs + self.warmup_epochs = warmup_epochs + self.model_lr = model_lr + self.model_wd = model_wd + self.arch_lr = arch_lr + self.stru_lr = stru_lr + self.lamb = lamb + + def train_stru(self, model, optimizer, data): + # forward + model.train() + data[0].adj = self.adjs + logits = model(data[0]).detach() + loss = 0 + for adj in self.adjs: + e1 = adj[0][0] + e2 = adj[0][1] + ew = adj[1] + diff = (logits[e1] - logits[e2]).pow(2).sum(1) + smooth = (diff * torch.sigmoid(ew)).sum() + dist = (ew * ew).sum() + #print(smooth,dist) + loss += self.lamb * smooth + dist + + optimizer.zero_grad() + loss.backward() + optimizer.step() + train_loss = loss.item() + del logits + + def _infer(self, model: BaseSpace, dataset, estimator: BaseEstimator, mask="train"): + dataset[0].adj = self.adjs + metric, loss = estimator.infer(model, dataset, mask=mask) + return metric, loss + + def prepare(self, dset): + """Train Pro-GNN. + """ + data = dset[0] + self.ews = [] + self.edges = data.edge_index.to(self.device) + edge_weight = torch.ones(self.edges.size(1)).to(self.device) + + self.adjs = [] + for i in range(self.steps): + edge_weight = Variable(edge_weight * 1.0, requires_grad = True).to(self.device) + self.ews.append(edge_weight) + self.adjs.append((self.edges, edge_weight)) + + def fit(self, data): + self.optimizer = optim.Adam(self.space.parameters(), lr=self.model_lr, weight_decay=self.model_wd) + self.arch_optimizer = optim.Adam(self.space.arch_parameters(), + lr=self.arch_lr, betas=(0.5, 0.999)) + self.stru_optimizer = optim.SGD(self.ews, lr=self.stru_lr) + + # Train model + best_performance = 0 + min_val_loss = float("inf") + min_train_loss = float("inf") + + t_total = time.time() + for epoch in range(self.num_epochs): + self.optimizer.zero_grad() + acc1, loss = self._infer(self.space, data, self.estimator, "train") + loss.backward() + self.optimizer.step() + + if epoch <20: + continue + self.train_stru(self.space, self.stru_optimizer, data) + + self.arch_optimizer.zero_grad() + acc2, loss = self._infer(self.space, data, self.estimator, "train") + loss.backward() + self.arch_optimizer.step() + + val_acc, val_loss = self._infer(self.space, data, self.estimator, "val") + if val_loss < min_val_loss: + min_val_loss = val_loss + best_performance = val_acc + self.space.keep_prediction() + + return best_performance, min_val_loss + + def search(self, space: BaseSpace, dataset, estimator): + self.estimator = estimator + self.space = space.to(self.device) + self.steps = space.steps + self.prepare(dataset) + perf, val_loss = self.fit(dataset) + return space.parse_model(None, self.device) \ No newline at end of file diff --git a/autogl/module/nas/space/__init__.py b/autogl/module/nas/space/__init__.py index 6d1e8f1..0a59ecd 100644 --- a/autogl/module/nas/space/__init__.py +++ b/autogl/module/nas/space/__init__.py @@ -23,6 +23,9 @@ from .graph_nas_macro import GraphNasMacroNodeClassificationSpace from .graph_nas import GraphNasNodeClassificationSpace from .single_path import SinglePathNodeClassificationSpace +from ..backend import * +if not is_dgl(): + from .gasso import GassoSpace def build_nas_space_from_name(name: str) -> BaseSpace: """ @@ -51,3 +54,6 @@ __all__ = [ "GraphNasNodeClassificationSpace", "SinglePathNodeClassificationSpace", ] + +if not is_dgl(): + __all__.append("GassoSpace") diff --git a/autogl/module/nas/space/gasso.py b/autogl/module/nas/space/gasso.py new file mode 100644 index 0000000..5a16fe7 --- /dev/null +++ b/autogl/module/nas/space/gasso.py @@ -0,0 +1,308 @@ +import typing as _typ + +from . import register_nas_space +from .base import apply_fixed_architecture +from .base import BaseSpace +from ...model import BaseModel +from ....utils import get_logger + +from ..backend import * +from ..utils import count_parameters, measure_latency + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn import Module +from .gasso_space import * +from torch.autograd import Variable +from collections import namedtuple + +Genotype = namedtuple('Genotype', 'normal normal_concat') +Genotype_normal = namedtuple('Genotype_normal', 'normal normal_concat') + +gnn_list = [ + "gat", # GAT with 2 heads + "gcn", # GCN + "gin", # GIN + #"cheb", # chebnet + "sage", # sage + #"arma", + #"sg", # simplifying gcn + "linear", # skip connection + #"skip", # skip connection + #"zero", # skip connection +] +act_list = [ + "sigmoid", "tanh", "relu", "linear", "elu" +] + +def gnn_map(gnn_name, in_dim, out_dim, concat=False, bias=True) -> Module: + ''' + + :param gnn_name: + :param in_dim: + :param out_dim: + :param concat: for gat, concat multi-head output or not + :return: GNN model + ''' + norm= True + if gnn_name == "gat": + return GATConv(in_dim, out_dim, 1, bias=bias, concat = False, add_self_loops=norm) + elif gnn_name == "gcn": + return GCNConv(in_dim, out_dim, add_self_loops=True, normalize=norm) + elif gnn_name == "gin": + return GINConv(torch.nn.Linear(in_dim, out_dim)) + elif gnn_name == "cheb": + return ChebConv(in_dim, out_dim, K=2, bias=bias) + elif gnn_name == "sage": + return SAGEConv(in_dim, out_dim, bias=bias) + elif gnn_name == "gated": + return GatedGraphConv(in_dim, out_dim, bias=bias) + elif gnn_name == "arma": + return ARMAConv(in_dim, out_dim, bias=bias, normalize=norm) + elif gnn_name == "sg": + return SGConv(in_dim, out_dim, bias=bias, normalize=norm) + elif gnn_name == "linear": + return LinearConv(in_dim, out_dim, bias=bias) + elif gnn_name == "skip": + return SkipConv(in_dim, out_dim, bias=bias) + elif gnn_name == "zero": + return ZeroConv(in_dim, out_dim, bias=bias) + else: + print(gnn_name) + +def Get_edges(adjs, ): + edges = [] + edges_weights = [] + for adj in adjs: + edges.append(adj[0]) + edges_weights.append(torch.sigmoid(adj[1])) + return edges, edges_weights + +class LinearConv(Module): + def __init__(self, + in_channels, + out_channels, + bias=True): + super(LinearConv, self).__init__() + + self.in_channels = in_channels + self.out_channels = out_channels + self.linear = torch.nn.Linear(in_channels, out_channels, bias) + + def forward(self, x, edge_index, edge_weight=None): + return self.linear(x) + + def __repr__(self): + return '{}({}, {})'.format(self.__class__.__name__, self.in_channels, + self.out_channels) + +class SkipConv(Module): + def __init__(self, + in_channels, + out_channels, + bias=True): + super(SkipConv, self).__init__() + self.out_dim = out_channels + + + def forward(self, x, edge_index, edge_weight=None): + return x + + def __repr__(self): + return '{}({}, {})'.format(self.__class__.__name__, self.in_channels, + self.out_channels) + +class ZeroConv(Module): + def __init__(self, + in_channels, + out_channels, + bias=True): + super(ZeroConv, self).__init__() + self.out_dim = out_channels + + + def forward(self, x, edge_index, edge_weight=None): + return torch.zeros([x.size(0), self.out_dim]).to(x.device) + + def __repr__(self): + return '{}({}, {})'.format(self.__class__.__name__, self.in_channels, + self.out_channels) + +class MixedOp(nn.Module): + + def __init__(self, in_c, out_c): + super(MixedOp, self).__init__() + self._ops = nn.ModuleList() + for action in gnn_list: + self._ops.append(gnn_map(action, in_c, out_c)) + + def forward(self, x, edge_index, edge_weight, weights, selected_idx=None): + if selected_idx is None: + fin = [] + for w, op, op_name in zip(weights, self._ops, gnn_list): + """if op_name == "gcn": + w = 1.0 + else: + continue""" + if edge_weight == None: + fin.append(w * op(x, edge_index)) + else: + fin.append(w * op(x, edge_index, edge_weight = edge_weight)) + return sum(fin) + #return sum(w * op(x, edge_index) for w, op in zip(weights, self._ops)) + else: # unchosen operations are pruned + return self._ops[selected_idx](x, edge_index) + +class CellWS(nn.Module): + + def __init__(self, steps, his_dim, hidden_dim, out_dim, dp, bias=True): + super(CellWS, self).__init__() + self.steps = steps + self._ops = nn.ModuleList() + self._bns = nn.ModuleList() + self.use2 = False + self.dp = 0.8 + for i in range(self.steps): + if i == 0: + inpdim = his_dim + else: + inpdim = hidden_dim + if i == self.steps - 1: + oupdim = out_dim + else: + oupdim = hidden_dim + op = MixedOp(inpdim, oupdim) + self._ops.append(op) + self._bns.append(nn.BatchNorm1d(oupdim)) + + def forward(self, x, adjs, weights): + edges, ews = Get_edges(adjs) + for i in range(self.steps): + if i > 0: + x = F.relu(x) + x = F.dropout(x, p=self.dp, training=self.training) + x = self._ops[i](x, edges[i], ews[i], weights[i]) # call the gcn module + return x + +@register_nas_space("gassospace") +class GassoSpace(BaseSpace): + def __init__( + self, + hidden_dim: _typ.Optional[int] = 64, + layer_number: _typ.Optional[int] = 2, + dropout: _typ.Optional[float] = 0.2, + input_dim: _typ.Optional[int] = None, + output_dim: _typ.Optional[int] = None, + ops: _typ.Tuple = gnn_list, + ): + super().__init__() + self.input_dim = input_dim + self.output_dim = output_dim + self.hidden_dim = hidden_dim + self.steps = layer_number + self.dropout = dropout + self.ops = ops + self.use_forward = True + self.dead_tensor = torch.nn.Parameter(torch.FloatTensor([1]), requires_grad = True) + + def instantiate( + self, + hidden_dim: _typ.Optional[int] = 64, + layer_number: _typ.Optional[int] = 2, + dropout: _typ.Optional[float] = 0.2, + input_dim: _typ.Optional[int] = None, + output_dim: _typ.Optional[int] = None, + ops: _typ.Tuple = gnn_list, + ): + super().instantiate() + self.input_dim = input_dim or self.input_dim + self.output_dim = output_dim or self.output_dim + self.hidden_dim = hidden_dim or self.hidden_dim + self.steps = layer_number or self.steps + self.dropout = dropout or self.dropout + self.ops = ops or self.ops + his_dim, cur_dim, hidden_dim, out_dim = self.input_dim, self.input_dim, self.hidden_dim, self.hidden_dim + self.cells = nn.ModuleList() + + self.cell = CellWS(self.steps, his_dim, hidden_dim, self.output_dim, self.dropout) + his_dim = cur_dim + cur_dim = self.steps * out_dim + + self.classifier = nn.Linear(cur_dim, self.output_dim) + + self.initialize_alphas() + + def forward(self, data): + if self.use_forward: + x, adjs = data.x, data.adj + x = F.dropout(x, p=self.dropout, training=self.training) + + weights = [] + for j in range(self.steps): + weights.append(F.softmax(self.alphas_normal[j], dim=-1)) + + x = self.cell(x, adjs, weights) + x = F.log_softmax(x, dim=1) + self.current_pred = x.detach() + return x + else: + #for i in self.parameters(): + # print(i) + x = self.prediction + self.dead_tensor * 0 + return x + + def keep_prediction(self): + self.prediction = self.current_pred + + def initialize_alphas(self): + num_ops = len(self.ops) + + self.alphas_normal = [] + for i in range(self.steps): + self.alphas_normal.append(Variable(1e-3 * torch.randn(num_ops).cuda(), requires_grad=True)) + + self._arch_parameters = [ + self.alphas_normal + ] + + def arch_parameters(self): + return self.alphas_normal + + def parse_gene_force(self, alphas): + gene = [] + n = 2 + max_num_edges = 1 + start = 0 + mat = F.softmax(torch.stack(alphas, dim=0), dim=-1).detach() + importance = torch.sum(mat[:, 1:], dim=-1).cpu() + for i in range(self.steps): + #end = start + 2 + end = start + n + num_edges_to_select = max_num_edges + if num_edges_to_select > 0: + post_select_edges = torch.topk(importance[start: end], k=num_edges_to_select).indices + start + else: + post_select_edges = [] + for j in range(start, end): + if num_edges_to_select <= 0: + raise Exception("Unknown errors") + else: + if j in post_select_edges: + idx = torch.argmax(alphas[j][1:]) + 1 + gene.append((self.ops[idx], j - start)) + start = end + n += 1 + + return gene + + def get_genotype(self, force=True): + gene_normal = self.parse_gene_force(self.alphas_normal) + n = 2 + concat = range(n, self.steps + n) + genotype = Genotype(normal=gene_normal, normal_concat=concat) + return genotype + + def parse_model(self, selection, device) -> BaseModel: + self.use_forward = False + return self.wrap(device) diff --git a/autogl/module/nas/space/gasso_space/__init__.py b/autogl/module/nas/space/gasso_space/__init__.py new file mode 100644 index 0000000..6d5d99a --- /dev/null +++ b/autogl/module/nas/space/gasso_space/__init__.py @@ -0,0 +1,23 @@ +from .message_passing import MessagePassing +from .gcn_conv import GCNConv +from .cheb_conv import ChebConv +from .sage_conv import SAGEConv +from .gat_conv import GATConv +from .gin_conv import GINConv, GINEConv +from .arma_conv import ARMAConv +from .edge_conv import EdgeConv, DynamicEdgeConv + +__all__ = [ + 'MessagePassing', + 'GCNConv', + 'ChebConv', + 'SAGEConv', + 'GATConv', + 'GINConv', + 'GINEConv', + 'ARMAConv', + 'EdgeConv', + 'DynamicEdgeConv', +] + +classes = __all__ diff --git a/autogl/module/nas/space/gasso_space/arma_conv.py b/autogl/module/nas/space/gasso_space/arma_conv.py new file mode 100644 index 0000000..12041ec --- /dev/null +++ b/autogl/module/nas/space/gasso_space/arma_conv.py @@ -0,0 +1,133 @@ +from typing import Callable +from torch_geometric.typing import Adj, OptTensor + +import torch +from torch import Tensor +from torch.nn import Parameter +import torch.nn.functional as F +from torch_sparse import SparseTensor, matmul +from torch_geometric.nn.conv import MessagePassing +from torch_geometric.nn.conv.gcn_conv import gcn_norm + +from .inits import glorot, zeros + + +class ARMAConv(MessagePassing): + r"""The ARMA graph convolutional operator from the `"Graph Neural Networks + with Convolutional ARMA Filters" `_ paper + + .. math:: + \mathbf{X}^{\prime} = \frac{1}{K} \sum_{k=1}^K \mathbf{X}_k^{(T)}, + + with :math:`\mathbf{X}_k^{(T)}` being recursively defined by + + .. math:: + \mathbf{X}_k^{(t+1)} = \sigma \left( \mathbf{\hat{L}} + \mathbf{X}_k^{(t)} \mathbf{W} + \mathbf{X}^{(0)} \mathbf{V} \right), + + where :math:`\mathbf{\hat{L}} = \mathbf{I} - \mathbf{L} = \mathbf{D}^{-1/2} + \mathbf{A} \mathbf{D}^{-1/2}` denotes the + modified Laplacian :math:`\mathbf{L} = \mathbf{I} - \mathbf{D}^{-1/2} + \mathbf{A} \mathbf{D}^{-1/2}`. + + Args: + in_channels (int): Size of each input sample :math:`\mathbf{x}^{(t)}`. + out_channels (int): Size of each output sample + :math:`\mathbf{x}^{(t+1)}`. + num_stacks (int, optional): Number of parallel stacks :math:`K`. + (default: :obj:`1`). + num_layers (int, optional): Number of layers :math:`T`. + (default: :obj:`1`) + act (callable, optional): Activation function :math:`\sigma`. + (default: :meth:`torch.nn.functional.ReLU`) + shared_weights (int, optional): If set to :obj:`True` the layers in + each stack will share the same parameters. (default: :obj:`False`) + dropout (float, optional): Dropout probability of the skip connection. + (default: :obj:`0.`) + bias (bool, optional): If set to :obj:`False`, the layer will not learn + an additive bias. (default: :obj:`True`) + **kwargs (optional): Additional arguments of + :class:`torch_geometric.nn.conv.MessagePassing`. + """ + def __init__(self, in_channels: int, out_channels: int, + num_stacks: int = 1, num_layers: int = 1, + shared_weights: bool = False, act: Callable = F.relu, normalize = True, + dropout: float = 0., bias: bool = True, **kwargs): + kwargs.setdefault('aggr', 'add') + super(ARMAConv, self).__init__(**kwargs) + + self.in_channels = in_channels + self.out_channels = out_channels + self.num_stacks = num_stacks + self.num_layers = num_layers + self.act = act + self.shared_weights = shared_weights + self.dropout = dropout + self.normalize = normalize + + K, T, F_in, F_out = num_stacks, num_layers, in_channels, out_channels + T = 1 if self.shared_weights else T + + self.init_weight = Parameter(torch.Tensor(K, F_in, F_out)) + self.weight = Parameter(torch.Tensor(max(1, T - 1), K, F_out, F_out)) + self.root_weight = Parameter(torch.Tensor(T, K, F_in, F_out)) + + if bias: + self.bias = Parameter(torch.Tensor(T, K, 1, F_out)) + else: + self.register_parameter('bias', None) + + self.reset_parameters() + + def reset_parameters(self): + glorot(self.init_weight) + glorot(self.weight) + glorot(self.root_weight) + zeros(self.bias) + + def forward(self, x: Tensor, edge_index: Adj, + edge_weight: OptTensor = None) -> Tensor: + """""" + + if isinstance(edge_index, Tensor) and self.normalize: + edge_index, edge_weight = gcn_norm( # yapf: disable + edge_index, edge_weight, x.size(self.node_dim), + add_self_loops=False, dtype=x.dtype) + + elif isinstance(edge_index, SparseTensor) and self.normalize: + edge_index = gcn_norm( # yapf: disable + edge_index, edge_weight, x.size(self.node_dim), + add_self_loops=False, dtype=x.dtype) + + x = x.unsqueeze(-3) + out = x + for t in range(self.num_layers): + if t == 0: + out = out @ self.init_weight + else: + out = out @ self.weight[0 if self.shared_weights else t - 1] + + # propagate_type: (x: Tensor, edge_weight: OptTensor) + out = self.propagate(edge_index, x=out, edge_weight=edge_weight, + size=None) + + root = F.dropout(x, p=self.dropout, training=self.training) + out += root @ self.root_weight[0 if self.shared_weights else t] + + if self.bias is not None: + out += self.bias[0 if self.shared_weights else t] + + out = self.act(out) + + return out.mean(dim=-3) + + def message(self, x_j: Tensor, edge_weight: Tensor) -> Tensor: + return edge_weight.view(-1, 1) * x_j + + def message_and_aggregate(self, adj_t: SparseTensor, x: Tensor) -> Tensor: + return matmul(adj_t, x, reduce=self.aggr) + + def __repr__(self): + return '{}({}, {}, num_stacks={}, num_layers={})'.format( + self.__class__.__name__, self.in_channels, self.out_channels, + self.num_stacks, self.num_layers) diff --git a/autogl/module/nas/space/gasso_space/cheb_conv.py b/autogl/module/nas/space/gasso_space/cheb_conv.py new file mode 100644 index 0000000..5642564 --- /dev/null +++ b/autogl/module/nas/space/gasso_space/cheb_conv.py @@ -0,0 +1,157 @@ +from typing import Optional +from torch_geometric.typing import OptTensor + +import torch +from torch.nn import Parameter +from torch_geometric.nn.conv import MessagePassing +from torch_geometric.utils import remove_self_loops, add_self_loops +from torch_geometric.utils import get_laplacian + +from .inits import glorot, zeros + + +class ChebConv(MessagePassing): + r"""The chebyshev spectral graph convolutional operator from the + `"Convolutional Neural Networks on Graphs with Fast Localized Spectral + Filtering" `_ paper + + .. math:: + \mathbf{X}^{\prime} = \sum_{k=1}^{K} \mathbf{Z}^{(k)} \cdot + \mathbf{\Theta}^{(k)} + + where :math:`\mathbf{Z}^{(k)}` is computed recursively by + + .. math:: + \mathbf{Z}^{(1)} &= \mathbf{X} + + \mathbf{Z}^{(2)} &= \mathbf{\hat{L}} \cdot \mathbf{X} + + \mathbf{Z}^{(k)} &= 2 \cdot \mathbf{\hat{L}} \cdot + \mathbf{Z}^{(k-1)} - \mathbf{Z}^{(k-2)} + + and :math:`\mathbf{\hat{L}}` denotes the scaled and normalized Laplacian + :math:`\frac{2\mathbf{L}}{\lambda_{\max}} - \mathbf{I}`. + + Args: + in_channels (int): Size of each input sample. + out_channels (int): Size of each output sample. + K (int): Chebyshev filter size :math:`K`. + normalization (str, optional): The normalization scheme for the graph + Laplacian (default: :obj:`"sym"`): + + 1. :obj:`None`: No normalization + :math:`\mathbf{L} = \mathbf{D} - \mathbf{A}` + + 2. :obj:`"sym"`: Symmetric normalization + :math:`\mathbf{L} = \mathbf{I} - \mathbf{D}^{-1/2} \mathbf{A} + \mathbf{D}^{-1/2}` + + 3. :obj:`"rw"`: Random-walk normalization + :math:`\mathbf{L} = \mathbf{I} - \mathbf{D}^{-1} \mathbf{A}` + + You need to pass :obj:`lambda_max` to the :meth:`forward` method of + this operator in case the normalization is non-symmetric. + :obj:`\lambda_max` should be a :class:`torch.Tensor` of size + :obj:`[num_graphs]` in a mini-batch scenario and a + scalar/zero-dimensional tensor when operating on single graphs. + You can pre-compute :obj:`lambda_max` via the + :class:`torch_geometric.transforms.LaplacianLambdaMax` transform. + bias (bool, optional): If set to :obj:`False`, the layer will not learn + an additive bias. (default: :obj:`True`) + **kwargs (optional): Additional arguments of + :class:`torch_geometric.nn.conv.MessagePassing`. + """ + def __init__(self, in_channels, out_channels, K, normalization='sym', + bias=True, **kwargs): + kwargs.setdefault('aggr', 'add') + super(ChebConv, self).__init__(**kwargs) + + assert K > 0 + assert normalization in [None, 'sym', 'rw'], 'Invalid normalization' + + self.in_channels = in_channels + self.out_channels = out_channels + self.normalization = normalization + self.weight = Parameter(torch.Tensor(K, in_channels, out_channels)) + + if bias: + self.bias = Parameter(torch.Tensor(out_channels)) + else: + self.register_parameter('bias', None) + + self.reset_parameters() + + def reset_parameters(self): + glorot(self.weight) + zeros(self.bias) + + def __norm__(self, edge_index, num_nodes: Optional[int], + edge_weight: OptTensor, normalization: Optional[str], + lambda_max, dtype: Optional[int] = None, + batch: OptTensor = None): + + edge_index, edge_weight = remove_self_loops(edge_index, edge_weight) + + edge_index, edge_weight = get_laplacian(edge_index, edge_weight, + normalization, dtype, + num_nodes) + + if batch is not None and lambda_max.numel() > 1: + lambda_max = lambda_max[batch[edge_index[0]]] + + edge_weight = (2.0 * edge_weight) / lambda_max + edge_weight.masked_fill_(edge_weight == float('inf'), 0) + + edge_index, edge_weight = add_self_loops(edge_index, edge_weight, + fill_value=-1., + num_nodes=num_nodes) + assert edge_weight is not None + + return edge_index, edge_weight + + def forward(self, x, edge_index, edge_weight: OptTensor = None, + batch: OptTensor = None, lambda_max: OptTensor = None): + """""" + if self.normalization != 'sym' and lambda_max is None: + raise ValueError('You need to pass `lambda_max` to `forward() in`' + 'case the normalization is non-symmetric.') + + if lambda_max is None: + lambda_max = torch.tensor(2.0, dtype=x.dtype, device=x.device) + if not isinstance(lambda_max, torch.Tensor): + lambda_max = torch.tensor(lambda_max, dtype=x.dtype, + device=x.device) + assert lambda_max is not None + + edge_index, norm = self.__norm__(edge_index, x.size(self.node_dim), + edge_weight, self.normalization, + lambda_max, dtype=x.dtype, + batch=batch) + + Tx_0 = x + Tx_1 = x # Dummy. + out = torch.matmul(Tx_0, self.weight[0]) + + # propagate_type: (x: Tensor, norm: Tensor) + if self.weight.size(0) > 1: + Tx_1 = self.propagate(edge_index, x=x, norm=norm, size=None) + out = out + torch.matmul(Tx_1, self.weight[1]) + + for k in range(2, self.weight.size(0)): + Tx_2 = self.propagate(edge_index, x=Tx_1, norm=norm, size=None) + Tx_2 = 2. * Tx_2 - Tx_0 + out = out + torch.matmul(Tx_2, self.weight[k]) + Tx_0, Tx_1 = Tx_1, Tx_2 + + if self.bias is not None: + out += self.bias + + return out + + def message(self, x_j, norm): + return norm.view(-1, 1) * x_j + + def __repr__(self): + return '{}({}, {}, K={}, normalization={})'.format( + self.__class__.__name__, self.in_channels, self.out_channels, + self.weight.size(0), self.normalization) diff --git a/autogl/module/nas/space/gasso_space/edge_conv.py b/autogl/module/nas/space/gasso_space/edge_conv.py new file mode 100644 index 0000000..6c7ab22 --- /dev/null +++ b/autogl/module/nas/space/gasso_space/edge_conv.py @@ -0,0 +1,123 @@ +from typing import Callable, Union, Optional +from torch_geometric.typing import OptTensor, PairTensor, PairOptTensor, Adj + +import torch +from torch import Tensor +from torch_geometric.nn.conv import MessagePassing + +from .inits import reset + +try: + from torch_cluster import knn +except ImportError: + knn = None + + +class EdgeConv(MessagePassing): + r"""The edge convolutional operator from the `"Dynamic Graph CNN for + Learning on Point Clouds" `_ paper + + .. math:: + \mathbf{x}^{\prime}_i = \sum_{j \in \mathcal{N}(i)} + h_{\mathbf{\Theta}}(\mathbf{x}_i \, \Vert \, + \mathbf{x}_j - \mathbf{x}_i), + + where :math:`h_{\mathbf{\Theta}}` denotes a neural network, *.i.e.* a MLP. + + Args: + nn (torch.nn.Module): A neural network :math:`h_{\mathbf{\Theta}}` that + maps pair-wise concatenated node features :obj:`x` of shape + :obj:`[-1, 2 * in_channels]` to shape :obj:`[-1, out_channels]`, + *e.g.*, defined by :class:`torch.nn.Sequential`. + aggr (string, optional): The aggregation scheme to use + (:obj:`"add"`, :obj:`"mean"`, :obj:`"max"`). + (default: :obj:`"max"`) + **kwargs (optional): Additional arguments of + :class:`torch_geometric.nn.conv.MessagePassing`. + """ + def __init__(self, nn: Callable, aggr: str = 'max', **kwargs): + super(EdgeConv, self).__init__(aggr=aggr, **kwargs) + self.nn = nn + self.reset_parameters() + + def reset_parameters(self): + reset(self.nn) + + def forward(self, x: Union[Tensor, PairTensor], edge_index: Adj) -> Tensor: + """""" + if isinstance(x, Tensor): + x: PairTensor = (x, x) + # propagate_type: (x: PairTensor) + return self.propagate(edge_index, x=x, size=None) + + def message(self, x_i: Tensor, x_j: Tensor) -> Tensor: + return self.nn(torch.cat([x_i, x_j - x_i], dim=-1)) + + def __repr__(self): + return '{}(nn={})'.format(self.__class__.__name__, self.nn) + + +class DynamicEdgeConv(MessagePassing): + r"""The dynamic edge convolutional operator from the `"Dynamic Graph CNN + for Learning on Point Clouds" `_ paper + (see :class:`torch_geometric.nn.conv.EdgeConv`), where the graph is + dynamically constructed using nearest neighbors in the feature space. + + Args: + nn (torch.nn.Module): A neural network :math:`h_{\mathbf{\Theta}}` that + maps pair-wise concatenated node features :obj:`x` of shape + `:obj:`[-1, 2 * in_channels]` to shape :obj:`[-1, out_channels]`, + *e.g.* defined by :class:`torch.nn.Sequential`. + k (int): Number of nearest neighbors. + aggr (string): The aggregation operator to use (:obj:`"add"`, + :obj:`"mean"`, :obj:`"max"`). (default: :obj:`"max"`) + num_workers (int): Number of workers to use for k-NN computation. + Has no effect in case :obj:`batch` is not :obj:`None`, or the input + lies on the GPU. (default: :obj:`1`) + **kwargs (optional): Additional arguments of + :class:`torch_geometric.nn.conv.MessagePassing`. + """ + def __init__(self, nn: Callable, k: int, aggr: str = 'max', + num_workers: int = 1, **kwargs): + super(DynamicEdgeConv, + self).__init__(aggr=aggr, flow='target_to_source', **kwargs) + + if knn is None: + raise ImportError('`DynamicEdgeConv` requires `torch-cluster`.') + + self.nn = nn + self.k = k + self.num_workers = num_workers + self.reset_parameters() + + def reset_parameters(self): + reset(self.nn) + + def forward( + self, x: Union[Tensor, PairTensor], + batch: Union[OptTensor, Optional[PairTensor]] = None) -> Tensor: + """""" + if isinstance(x, Tensor): + x: PairTensor = (x, x) + assert x[0].dim() == 2, \ + 'Static graphs not supported in `DynamicEdgeConv`.' + + b: PairOptTensor = (None, None) + if isinstance(batch, Tensor): + b = (batch, batch) + elif isinstance(batch, tuple): + assert batch is not None + b = (batch[0], batch[1]) + + edge_index = knn(x[0], x[1], self.k, b[0], b[1], + num_workers=self.num_workers) + + # propagate_type: (x: PairTensor) + return self.propagate(edge_index, x=x, size=None) + + def message(self, x_i: Tensor, x_j: Tensor) -> Tensor: + return self.nn(torch.cat([x_i, x_j - x_i], dim=-1)) + + def __repr__(self): + return '{}(nn={}, k={})'.format(self.__class__.__name__, self.nn, + self.k) diff --git a/autogl/module/nas/space/gasso_space/gat_conv.py b/autogl/module/nas/space/gasso_space/gat_conv.py new file mode 100644 index 0000000..71586fe --- /dev/null +++ b/autogl/module/nas/space/gasso_space/gat_conv.py @@ -0,0 +1,194 @@ +from typing import Union, Tuple, Optional +from torch_geometric.typing import (OptPairTensor, Adj, Size, NoneType, + OptTensor) + +import torch +from torch import Tensor +import torch.nn.functional as F +from torch.nn import Parameter, Linear +from torch_sparse import SparseTensor, set_diag +from torch_geometric.nn.conv import MessagePassing +from torch_geometric.utils import remove_self_loops, add_self_loops, softmax + +from .inits import glorot, zeros + + +class GATConv(MessagePassing): + r"""The graph attentional operator from the `"Graph Attention Networks" + `_ paper + + .. math:: + \mathbf{x}^{\prime}_i = \alpha_{i,i}\mathbf{\Theta}\mathbf{x}_{i} + + \sum_{j \in \mathcal{N}(i)} \alpha_{i,j}\mathbf{\Theta}\mathbf{x}_{j}, + + where the attention coefficients :math:`\alpha_{i,j}` are computed as + + .. math:: + \alpha_{i,j} = + \frac{ + \exp\left(\mathrm{LeakyReLU}\left(\mathbf{a}^{\top} + [\mathbf{\Theta}\mathbf{x}_i \, \Vert \, \mathbf{\Theta}\mathbf{x}_j] + \right)\right)} + {\sum_{k \in \mathcal{N}(i) \cup \{ i \}} + \exp\left(\mathrm{LeakyReLU}\left(\mathbf{a}^{\top} + [\mathbf{\Theta}\mathbf{x}_i \, \Vert \, \mathbf{\Theta}\mathbf{x}_k] + \right)\right)}. + + Args: + in_channels (int or tuple): Size of each input sample. A tuple + corresponds to the sizes of source and target dimensionalities. + out_channels (int): Size of each output sample. + heads (int, optional): Number of multi-head-attentions. + (default: :obj:`1`) + concat (bool, optional): If set to :obj:`False`, the multi-head + attentions are averaged instead of concatenated. + (default: :obj:`True`) + negative_slope (float, optional): LeakyReLU angle of the negative + slope. (default: :obj:`0.2`) + dropout (float, optional): Dropout probability of the normalized + attention coefficients which exposes each node to a stochastically + sampled neighborhood during training. (default: :obj:`0`) + add_self_loops (bool, optional): If set to :obj:`False`, will not add + self-loops to the input graph. (default: :obj:`True`) + bias (bool, optional): If set to :obj:`False`, the layer will not learn + an additive bias. (default: :obj:`True`) + **kwargs (optional): Additional arguments of + :class:`torch_geometric.nn.conv.MessagePassing`. + """ + _alpha: OptTensor + + def __init__(self, in_channels: Union[int, Tuple[int, int]], + out_channels: int, heads: int = 1, concat: bool = True, + negative_slope: float = 0.2, dropout: float = 0., + add_self_loops: bool = True, bias: bool = True, **kwargs): + kwargs.setdefault('aggr', 'add') + super(GATConv, self).__init__(node_dim=0, **kwargs) + + self.in_channels = in_channels + self.out_channels = out_channels + self.heads = heads + self.concat = concat + self.negative_slope = negative_slope + self.dropout = dropout + self.add_self_loops = add_self_loops + + if isinstance(in_channels, int): + self.lin_l = Linear(in_channels, heads * out_channels, bias=False) + self.lin_r = self.lin_l + else: + self.lin_l = Linear(in_channels[0], heads * out_channels, False) + self.lin_r = Linear(in_channels[1], heads * out_channels, False) + + self.att_l = Parameter(torch.Tensor(1, heads, out_channels)) + self.att_r = Parameter(torch.Tensor(1, heads, out_channels)) + + if bias and concat: + self.bias = Parameter(torch.Tensor(heads * out_channels)) + elif bias and not concat: + self.bias = Parameter(torch.Tensor(out_channels)) + else: + self.register_parameter('bias', None) + + self._alpha = None + + self.reset_parameters() + + def reset_parameters(self): + glorot(self.lin_l.weight) + glorot(self.lin_r.weight) + glorot(self.att_l) + glorot(self.att_r) + zeros(self.bias) + + def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj, edge_weight: OptTensor = None, + size: Size = None, return_attention_weights=None): + # type: (Union[Tensor, OptPairTensor], Tensor, Size, NoneType) -> Tensor # noqa + # type: (Union[Tensor, OptPairTensor], SparseTensor, Size, NoneType) -> Tensor # noqa + # type: (Union[Tensor, OptPairTensor], Tensor, Size, bool) -> Tuple[Tensor, Tuple[Tensor, Tensor]] # noqa + # type: (Union[Tensor, OptPairTensor], SparseTensor, Size, bool) -> Tuple[Tensor, SparseTensor] # noqa + r""" + + Args: + return_attention_weights (bool, optional): If set to :obj:`True`, + will additionally return the tuple + :obj:`(edge_index, attention_weights)`, holding the computed + attention weights for each edge. (default: :obj:`None`) + """ + H, C = self.heads, self.out_channels + + x_l: OptTensor = None + x_r: OptTensor = None + alpha_l: OptTensor = None + alpha_r: OptTensor = None + if isinstance(x, Tensor): + assert x.dim() == 2, 'Static graphs not supported in `GATConv`.' + x_l = x_r = self.lin_l(x).view(-1, H, C) + alpha_l = (x_l * self.att_l).sum(dim=-1) + alpha_r = (x_r * self.att_r).sum(dim=-1) + else: + x_l, x_r = x[0], x[1] + assert x[0].dim() == 2, 'Static graphs not supported in `GATConv`.' + x_l = self.lin_l(x_l).view(-1, H, C) + alpha_l = (x_l * self.att_l).sum(dim=-1) + if x_r is not None: + x_r = self.lin_r(x_r).view(-1, H, C) + alpha_r = (x_r * self.att_r).sum(dim=-1) + + assert x_l is not None + assert alpha_l is not None + + if self.add_self_loops: + if isinstance(edge_index, Tensor): + num_nodes = x_l.size(0) + if x_r is not None: + num_nodes = min(num_nodes, x_r.size(0)) + if size is not None: + num_nodes = min(size[0], size[1]) + edge_index, edge_weight = remove_self_loops(edge_index, edge_attr=edge_weight) + if edge_weight != None: + edge_index, edge_weight = add_self_loops(edge_index, edge_weight=edge_weight, num_nodes=num_nodes) + else: + edge_index, _ = add_self_loops(edge_index, num_nodes=num_nodes) + elif isinstance(edge_index, SparseTensor): + edge_index = set_diag(edge_index) + + # propagate_type: (x: OptPairTensor, alpha: OptPairTensor) + out = self.propagate(edge_index, x=(x_l, x_r), + alpha=(alpha_l, alpha_r), edge_weight = edge_weight, size=size) + + alpha = self._alpha + self._alpha = None + + if self.concat: + out = out.view(-1, self.heads * self.out_channels) + else: + out = out.mean(dim=1) + + if self.bias is not None: + out += self.bias + + if isinstance(return_attention_weights, bool): + assert alpha is not None + if isinstance(edge_index, Tensor): + return out, (edge_index, alpha) + elif isinstance(edge_index, SparseTensor): + return out, edge_index.set_value(alpha, layout='coo') + else: + return out + + def message(self, x_j: Tensor, alpha_j: Tensor, alpha_i: OptTensor, + index: Tensor, ptr: OptTensor, + size_i: Optional[int], edge_weight: OptTensor = None) -> Tensor: + alpha = alpha_j if alpha_i is None else alpha_j + alpha_i + alpha = F.leaky_relu(alpha, self.negative_slope) + if edge_weight != None: + alpha = alpha.mul(edge_weight.unsqueeze(1)) + alpha = softmax(alpha, index, ptr, size_i) + self._alpha = alpha + alpha = F.dropout(alpha, p=self.dropout, training=self.training) + return x_j * alpha.unsqueeze(-1) + + def __repr__(self): + return '{}({}, {}, heads={})'.format(self.__class__.__name__, + self.in_channels, + self.out_channels, self.heads) diff --git a/autogl/module/nas/space/gasso_space/gcn_conv.py b/autogl/module/nas/space/gasso_space/gcn_conv.py new file mode 100644 index 0000000..3a41b9d --- /dev/null +++ b/autogl/module/nas/space/gasso_space/gcn_conv.py @@ -0,0 +1,200 @@ +from typing import Optional, Tuple +from torch_geometric.typing import Adj, OptTensor, PairTensor + +import torch +from torch import Tensor +from torch.nn import Parameter +from torch_scatter import scatter_add +from torch_sparse import SparseTensor, matmul, fill_diag, sum, mul +from torch_geometric.nn.conv import MessagePassing +from torch_geometric.utils import add_remaining_self_loops +from torch_geometric.utils.num_nodes import maybe_num_nodes + +from .inits import glorot, zeros + + +@torch.jit._overload +def gcn_norm(edge_index, edge_weight=None, num_nodes=None, improved=False, + add_self_loops=True, dtype=None): + # type: (Tensor, OptTensor, Optional[int], bool, bool, Optional[int]) -> PairTensor # noqa + pass + + +@torch.jit._overload +def gcn_norm(edge_index, edge_weight=None, num_nodes=None, improved=False, + add_self_loops=True, dtype=None): + # type: (SparseTensor, OptTensor, Optional[int], bool, bool, Optional[int]) -> SparseTensor # noqa + pass + + +def gcn_norm(edge_index, edge_weight=None, num_nodes=None, improved=False, + add_self_loops=True, dtype=None): + + fill_value = 2. if improved else 1. + + if isinstance(edge_index, SparseTensor): + adj_t = edge_index + if not adj_t.has_value(): + adj_t = adj_t.fill_value(1., dtype=dtype) + if add_self_loops: + adj_t = fill_diag(adj_t, fill_value) + deg = sum(adj_t, dim=1) + deg_inv_sqrt = deg.pow_(-0.5) + deg_inv_sqrt.masked_fill_(deg_inv_sqrt == float('inf'), 0.) + adj_t = mul(adj_t, deg_inv_sqrt.view(-1, 1)) + adj_t = mul(adj_t, deg_inv_sqrt.view(1, -1)) + return adj_t + + else: + num_nodes = maybe_num_nodes(edge_index, num_nodes) + + if edge_weight is None: + edge_weight = torch.ones((edge_index.size(1), ), dtype=dtype, + device=edge_index.device) + + if add_self_loops: + edge_index, tmp_edge_weight = add_remaining_self_loops( + edge_index, edge_weight, fill_value, num_nodes) + assert tmp_edge_weight is not None + edge_weight = tmp_edge_weight + + row, col = edge_index[0], edge_index[1] + deg = scatter_add(edge_weight, col, dim=0, dim_size=num_nodes) + deg_inv_sqrt = deg.pow_(-0.5) + deg_inv_sqrt.masked_fill_(deg_inv_sqrt == float('inf'), 0) + return edge_index, deg_inv_sqrt[row] * edge_weight * deg_inv_sqrt[col] + + +class GCNConv(MessagePassing): + r"""The graph convolutional operator from the `"Semi-supervised + Classification with Graph Convolutional Networks" + `_ paper + + .. math:: + \mathbf{X}^{\prime} = \mathbf{\hat{D}}^{-1/2} \mathbf{\hat{A}} + \mathbf{\hat{D}}^{-1/2} \mathbf{X} \mathbf{\Theta}, + + where :math:`\mathbf{\hat{A}} = \mathbf{A} + \mathbf{I}` denotes the + adjacency matrix with inserted self-loops and + :math:`\hat{D}_{ii} = \sum_{j=0} \hat{A}_{ij}` its diagonal degree matrix. + The adjacency matrix can include other values than :obj:`1` representing + edge weights via the optional :obj:`edge_weight` tensor. + + Its node-wise formulation is given by: + + .. math:: + \mathbf{x}^{\prime}_i = \mathbf{\Theta} \sum_{j} + \frac{1}{\sqrt{\hat{d}_j \hat{d}_i}} \mathbf{x}_j + + with :math:`\hat{d}_i = 1 + \sum_{j \in \mathcal{N}(i)} e_{j,i}`, where + :math:`e_{j,i}` denotes the edge weight from source node :obj:`i` to target + node :obj:`j` (default: :obj:`1`) + + Args: + in_channels (int): Size of each input sample. + out_channels (int): Size of each output sample. + improved (bool, optional): If set to :obj:`True`, the layer computes + :math:`\mathbf{\hat{A}}` as :math:`\mathbf{A} + 2\mathbf{I}`. + (default: :obj:`False`) + cached (bool, optional): If set to :obj:`True`, the layer will cache + the computation of :math:`\mathbf{\hat{D}}^{-1/2} \mathbf{\hat{A}} + \mathbf{\hat{D}}^{-1/2}` on first execution, and will use the + cached version for further executions. + This parameter should only be set to :obj:`True` in transductive + learning scenarios. (default: :obj:`False`) + add_self_loops (bool, optional): If set to :obj:`False`, will not add + self-loops to the input graph. (default: :obj:`True`) + normalize (bool, optional): Whether to add self-loops and apply + symmetric normalization. (default: :obj:`True`) + bias (bool, optional): If set to :obj:`False`, the layer will not learn + an additive bias. (default: :obj:`True`) + **kwargs (optional): Additional arguments of + :class:`torch_geometric.nn.conv.MessagePassing`. + """ + + _cached_edge_index: Optional[Tuple[Tensor, Tensor]] + _cached_adj_t: Optional[SparseTensor] + + def __init__(self, in_channels: int, out_channels: int, + improved: bool = False, cached: bool = False, + add_self_loops: bool = True, normalize: bool = True, + bias: bool = True, **kwargs): + + kwargs.setdefault('aggr', 'add') + super(GCNConv, self).__init__(**kwargs) + + self.in_channels = in_channels + self.out_channels = out_channels + self.improved = improved + self.cached = cached + self.add_self_loops = add_self_loops + self.normalize = normalize + + self._cached_edge_index = None + self._cached_adj_t = None + + self.weight = Parameter(torch.Tensor(in_channels, out_channels)) + + if bias: + self.bias = Parameter(torch.Tensor(out_channels)) + else: + self.register_parameter('bias', None) + + self.reset_parameters() + + def reset_parameters(self): + glorot(self.weight) + zeros(self.bias) + self._cached_edge_index = None + self._cached_adj_t = None + + def forward(self, x: Tensor, edge_index: Adj, + edge_weight: OptTensor = None) -> Tensor: + """""" + + if self.normalize: + if isinstance(edge_index, Tensor): + cache = self._cached_edge_index + if cache is None: + edge_index, edge_weight = gcn_norm( # yapf: disable + edge_index, edge_weight, x.size(self.node_dim), + self.improved, self.add_self_loops, dtype=x.dtype) + if self.cached: + self._cached_edge_index = (edge_index, edge_weight) + else: + edge_index, edge_weight = cache[0], cache[1] + + elif isinstance(edge_index, SparseTensor): + cache = self._cached_adj_t + if cache is None: + edge_index = gcn_norm( # yapf: disable + edge_index, edge_weight, x.size(self.node_dim), + self.improved, self.add_self_loops, dtype=x.dtype) + if self.cached: + self._cached_adj_t = edge_index + else: + edge_index = cache + + x = torch.matmul(x, self.weight) + + # propagate_type: (x: Tensor, edge_weight: OptTensor) + out = self.propagate(edge_index, x=x, edge_weight=edge_weight, + size=None) + + if self.bias is not None: + out += self.bias + + return out + + def message(self, x_j: Tensor, edge_weight: OptTensor) -> Tensor: + if edge_weight is None: + return x_j + else: + return edge_weight.view(-1, 1) * x_j + + def message_and_aggregate(self, adj_t: SparseTensor, x: Tensor) -> Tensor: + return matmul(adj_t, x, reduce=self.aggr) + + def __repr__(self): + return '{}({}, {})'.format(self.__class__.__name__, self.in_channels, + self.out_channels) diff --git a/autogl/module/nas/space/gasso_space/gin_conv.py b/autogl/module/nas/space/gasso_space/gin_conv.py new file mode 100644 index 0000000..b7a3b2e --- /dev/null +++ b/autogl/module/nas/space/gasso_space/gin_conv.py @@ -0,0 +1,157 @@ +from typing import Callable, Union +from torch_geometric.typing import OptPairTensor, Adj, OptTensor, Size + +import torch +from torch import Tensor +import torch.nn.functional as F +from torch_sparse import SparseTensor, matmul +from torch_geometric.nn.conv import MessagePassing + +from .inits import reset + + +class GINConv(MessagePassing): + r"""The graph isomorphism operator from the `"How Powerful are + Graph Neural Networks?" `_ paper + + .. math:: + \mathbf{x}^{\prime}_i = h_{\mathbf{\Theta}} \left( (1 + \epsilon) \cdot + \mathbf{x}_i + \sum_{j \in \mathcal{N}(i)} \mathbf{x}_j \right) + + or + + .. math:: + \mathbf{X}^{\prime} = h_{\mathbf{\Theta}} \left( \left( \mathbf{A} + + (1 + \epsilon) \cdot \mathbf{I} \right) \cdot \mathbf{X} \right), + + here :math:`h_{\mathbf{\Theta}}` denotes a neural network, *.i.e.* an MLP. + + Args: + nn (torch.nn.Module): A neural network :math:`h_{\mathbf{\Theta}}` that + maps node features :obj:`x` of shape :obj:`[-1, in_channels]` to + shape :obj:`[-1, out_channels]`, *e.g.*, defined by + :class:`torch.nn.Sequential`. + eps (float, optional): (Initial) :math:`\epsilon`-value. + (default: :obj:`0.`) + train_eps (bool, optional): If set to :obj:`True`, :math:`\epsilon` + will be a trainable parameter. (default: :obj:`False`) + **kwargs (optional): Additional arguments of + :class:`torch_geometric.nn.conv.MessagePassing`. + """ + def __init__(self, nn: Callable, eps: float = 0., train_eps: bool = False, + **kwargs): + kwargs.setdefault('aggr', 'add') + super(GINConv, self).__init__(**kwargs) + self.nn = nn + self.initial_eps = eps + if train_eps: + self.eps = torch.nn.Parameter(torch.Tensor([eps])) + else: + self.register_buffer('eps', torch.Tensor([eps])) + self.reset_parameters() + + def reset_parameters(self): + reset(self.nn) + self.eps.data.fill_(self.initial_eps) + + def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj, edge_weight: OptTensor = None, + size: Size = None) -> Tensor: + """""" + if isinstance(x, Tensor): + x: OptPairTensor = (x, x) + + # propagate_type: (x: OptPairTensor) + out = self.propagate(edge_index, x=x, edge_weight = edge_weight, size=size) + + x_r = x[1] + if x_r is not None: + out += (1 + self.eps) * x_r + + return self.nn(out) + + #def message(self, x_j: Tensor) -> Tensor: + # return x_j + def message(self, x_j: Tensor, edge_weight: OptTensor) -> Tensor: + if edge_weight is None: + return x_j + else: + return edge_weight.view(-1, 1) * x_j + + + def message_and_aggregate(self, adj_t: SparseTensor, + x: OptPairTensor) -> Tensor: + adj_t = adj_t.set_value(None, layout=None) + return matmul(adj_t, x[0], reduce=self.aggr) + + def __repr__(self): + return '{}(nn={})'.format(self.__class__.__name__, self.nn) + + +class GINEConv(MessagePassing): + r"""The modified :class:`GINConv` operator from the `"Strategies for + Pre-training Graph Neural Networks" `_ + paper + + .. math:: + \mathbf{x}^{\prime}_i = h_{\mathbf{\Theta}} \left( (1 + \epsilon) \cdot + \mathbf{x}_i + \sum_{j \in \mathcal{N}(i)} \mathrm{ReLU} + ( \mathbf{x}_j + \mathbf{e}_{j,i} ) \right) + + that is able to incorporate edge features :math:`\mathbf{e}_{j,i}` into + the aggregation procedure. + + Args: + nn (torch.nn.Module): A neural network :math:`h_{\mathbf{\Theta}}` that + maps node features :obj:`x` of shape :obj:`[-1, in_channels]` to + shape :obj:`[-1, out_channels]`, *e.g.*, defined by + :class:`torch.nn.Sequential`. + eps (float, optional): (Initial) :math:`\epsilon`-value. + (default: :obj:`0.`) + train_eps (bool, optional): If set to :obj:`True`, :math:`\epsilon` + will be a trainable parameter. (default: :obj:`False`) + **kwargs (optional): Additional arguments of + :class:`torch_geometric.nn.conv.MessagePassing`. + """ + def __init__(self, nn: Callable, eps: float = 0., train_eps: bool = False, + **kwargs): + kwargs.setdefault('aggr', 'add') + super(GINEConv, self).__init__(**kwargs) + self.nn = nn + self.initial_eps = eps + if train_eps: + self.eps = torch.nn.Parameter(torch.Tensor([eps])) + else: + self.register_buffer('eps', torch.Tensor([eps])) + self.reset_parameters() + + def reset_parameters(self): + reset(self.nn) + self.eps.data.fill_(self.initial_eps) + + def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj, + edge_attr: OptTensor = None, size: Size = None) -> Tensor: + """""" + if isinstance(x, Tensor): + x: OptPairTensor = (x, x) + + # Node and edge feature dimensionalites need to match. + if isinstance(edge_index, Tensor): + assert edge_attr is not None + assert x[0].size(-1) == edge_attr.size(-1) + elif isinstance(edge_index, SparseTensor): + assert x[0].size(-1) == edge_index.size(-1) + + # propagate_type: (x: OptPairTensor, edge_attr: OptTensor) + out = self.propagate(edge_index, x=x, edge_attr=edge_attr, size=size) + + x_r = x[1] + if x_r is not None: + out += (1 + self.eps) * x_r + + return self.nn(out) + + def message(self, x_j: Tensor, edge_attr: Tensor) -> Tensor: + return F.relu(x_j + edge_attr) + + def __repr__(self): + return '{}(nn={})'.format(self.__class__.__name__, self.nn) diff --git a/autogl/module/nas/space/gasso_space/inits.py b/autogl/module/nas/space/gasso_space/inits.py new file mode 100644 index 0000000..9a22541 --- /dev/null +++ b/autogl/module/nas/space/gasso_space/inits.py @@ -0,0 +1,56 @@ +import math + +import torch + + +def uniform(size, tensor): + if tensor is not None: + bound = 1.0 / math.sqrt(size) + tensor.data.uniform_(-bound, bound) + + +def kaiming_uniform(tensor, fan, a): + if tensor is not None: + bound = math.sqrt(6 / ((1 + a**2) * fan)) + tensor.data.uniform_(-bound, bound) + + +def glorot(tensor): + if tensor is not None: + stdv = math.sqrt(6.0 / (tensor.size(-2) + tensor.size(-1))) + tensor.data.uniform_(-stdv, stdv) + + +def glorot_orthogonal(tensor, scale): + if tensor is not None: + torch.nn.init.orthogonal_(tensor.data) + scale /= ((tensor.size(-2) + tensor.size(-1)) * tensor.var()) + tensor.data *= scale.sqrt() + + +def zeros(tensor): + if tensor is not None: + tensor.data.fill_(0) + + +def ones(tensor): + if tensor is not None: + tensor.data.fill_(1) + + +def normal(tensor, mean, std): + if tensor is not None: + tensor.data.normal_(mean, std) + + +def reset(nn): + def _reset(item): + if hasattr(item, 'reset_parameters'): + item.reset_parameters() + + if nn is not None: + if hasattr(nn, 'children') and len(list(nn.children())) > 0: + for item in nn.children(): + _reset(item) + else: + _reset(nn) diff --git a/autogl/module/nas/space/gasso_space/message_passing.jinja b/autogl/module/nas/space/gasso_space/message_passing.jinja new file mode 100644 index 0000000..4c1dacb --- /dev/null +++ b/autogl/module/nas/space/gasso_space/message_passing.jinja @@ -0,0 +1,153 @@ +from typing import * +from torch_geometric.typing import * + +import torch +from torch import Tensor +import torch_sparse +from torch_sparse import SparseTensor +from torch_geometric.nn.conv.message_passing import * +from {{module}} import * + + +class Propagate_{{uid}}(NamedTuple): +{%- for k, v in prop_types.items() %} + {{k}}: {{v}} +{%- endfor %} + + +class Collect_{{uid}}(NamedTuple): +{%- for k, v in collect_types.items() %} + {{k}}: {{v}} +{%- endfor %} + + + +class {{cls_name}}({{parent_cls_name}}): + + @torch.jit._overload_method + def __check_input__(self, edge_index, size): + # type: (Tensor, Size) -> List[Optional[int]] + pass + + @torch.jit._overload_method + def __check_input__(self, edge_index, size): + # type: (SparseTensor, Size) -> List[Optional[int]] + pass + +{{check_input}} + + @torch.jit._overload_method + def __lift__(self, src, edge_index, dim): + # type: (Tensor, Tensor, int) -> Tensor + pass + + @torch.jit._overload_method + def __lift__(self, src, edge_index, dim): + # type: (Tensor, SparseTensor, int) -> Tensor + pass + +{{lift}} + + @torch.jit._overload_method + def __collect__(self, edge_def, size, kwargs): + # type: (Tensor, List[Optional[int]], Propagate_{{uid}}) -> Collect_{{uid}} + pass + + @torch.jit._overload_method + def __collect__(self, edge_def, size, kwargs): + # type: (SparseTensor, List[Optional[int]], Propagate_{{uid}}) -> Collect_{{uid}} + pass + + def __collect__(self, edge_def, size, kwargs): + init = torch.tensor(0.) + i, j = (1, 0) if self.flow == 'source_to_target' else (0, 1) +{% for arg in user_args %} +{%- if arg[-2:] not in ['_i', '_j'] %} + {{arg}} = kwargs.{{arg}} +{%- else %} + {{arg}}: {{collect_types[arg]}} = {% if collect_types[arg][:8] == 'Optional' %}None{% else %}init{% endif %} + data = kwargs.{{arg[:-2]}} + if isinstance(data, (tuple, list)): + assert len(data) == 2 +{%- if arg[-2:] == '_j' %} + tmp = data[1] + if isinstance(tmp, Tensor): + self.__set_size__(size, 1, tmp) + {{arg}} = data[0] +{%- else %} + tmp = data[0] + if isinstance(tmp, Tensor): + self.__set_size__(size, 0, tmp) + {{arg}} = data[1] +{%- endif %} + else: + {{arg}} = data + if isinstance({{arg}}, Tensor): + self.__set_size__(size, {% if arg[-2:] == '_j'%}0{% else %}1{% endif %}, {{arg}}) + {{arg}} = self.__lift__({{arg}}, edge_def, {% if arg[-2:] == "_j" %}j{% else %}i{% endif %}) +{%- endif %} +{%- endfor %} + + edge_index: Optional[Tensor] = None + adj_t: Optional[SparseTensor] = None + edge_index_i: torch.Tensor = init + edge_index_j: torch.Tensor = init + ptr: Optional[Tensor] = None + if isinstance(edge_def, Tensor): + edge_index = edge_def + edge_index_i = edge_def[i] + edge_index_j = edge_def[j] + elif isinstance(edge_def, SparseTensor): + adj_t = edge_def + edge_index_i = edge_def.storage.row() + edge_index_j = edge_def.storage.col() + ptr = edge_def.storage.rowptr() + {% if 'edge_weight' in collect_types.keys() %}edge_weight = edge_def.storage.value(){% endif %} + {% if 'edge_attr' in collect_types.keys() %}edge_attr = edge_def.storage.value(){% endif %} + {% if 'edge_type' in collect_types.keys() %}edge_type = edge_def.storage.value(){% endif %} + + {% if collect_types.get('edge_weight', 'Optional')[:8] != 'Optional' %}assert edge_weight is not None{% endif %} + {% if collect_types.get('edge_attr', 'Optional')[:8] != 'Optional' %}assert edge_attr is not None{% endif %} + {% if collect_types.get('edge_type', 'Optional')[:8] != 'Optional' %}assert edge_type is not None{% endif %} + + index = edge_index_i + size_i = size[1] if size[1] is not None else size[0] + size_j = size[0] if size[0] is not None else size[1] + dim_size = size_i + + return Collect_{{uid}}({% for k in collect_types.keys() %}{{k}}={{k}}{{ ", " if not loop.last }}{% endfor %}) + + @torch.jit._overload_method + def propagate(self, edge_index, {{ prop_types.keys()|join(', ') }}, size=None): + # type: (Tensor, {{ prop_types.values()|join(', ') }}, Size) -> Tensor + pass + + @torch.jit._overload_method + def propagate(self, edge_index, {{ prop_types.keys()|join(', ') }}, size=None): + # type: (SparseTensor, {{ prop_types.values()|join(', ') }}, Size) -> Tensor + pass + + def propagate(self, edge_index, {{ prop_types.keys()|join(', ') }}, size=None): + the_size = self.__check_input__(edge_index, size) + in_kwargs = Propagate_{{uid}}({% for k in prop_types.keys() %}{{k}}={{k}}{{ ", " if not loop.last }}{% endfor %}) + + if self.fuse: + if isinstance(edge_index, SparseTensor): + out = self.message_and_aggregate(edge_index{% for k in msg_and_aggr_args %}, {{k}}=in_kwargs.{{k}}{% endfor %}) + return self.update(out{% for k in update_args %}, {{k}}=in_kwargs.{{k}}{% endfor %}) + + kwargs = self.__collect__(edge_index, the_size, in_kwargs) + out = self.message({% for k in msg_args %}{{k}}=kwargs.{{k}}{{ ", " if not loop.last }}{% endfor %}) + out = self.aggregate(out{% for k in aggr_args %}, {{k}}=kwargs.{{k}}{% endfor %}) + return self.update(out{% for k in update_args %}, {{k}}=kwargs.{{k}}{% endfor %}) + +{%- for (arg_types, return_type_repr) in forward_types %} + + @torch.jit._overload_method + {{forward_header}} + # type: ({{arg_types|join(', ')}}) -> {{return_type_repr}} + pass +{%- endfor %} + + {{forward_header}} +{{forward_body}} diff --git a/autogl/module/nas/space/gasso_space/message_passing.py b/autogl/module/nas/space/gasso_space/message_passing.py new file mode 100644 index 0000000..6e4a97d --- /dev/null +++ b/autogl/module/nas/space/gasso_space/message_passing.py @@ -0,0 +1,389 @@ +import os +import re +import inspect +import os.path as osp +from uuid import uuid1 +from itertools import chain +from inspect import Parameter +from typing import List, Optional, Set +from torch_geometric.typing import Adj, Size + +import torch +from torch import Tensor +from jinja2 import Template +from torch_sparse import SparseTensor +from torch_scatter import gather_csr, scatter, segment_csr + +from .utils.helpers import expand_left +from .utils.jit import class_from_module_repr +from .utils.typing import (sanitize, split_types_repr, parse_types, + resolve_types) +from .utils.inspector import Inspector, func_header_repr, func_body_repr + + +class MessagePassing(torch.nn.Module): + r"""Base class for creating message passing layers of the form + + .. math:: + \mathbf{x}_i^{\prime} = \gamma_{\mathbf{\Theta}} \left( \mathbf{x}_i, + \square_{j \in \mathcal{N}(i)} \, \phi_{\mathbf{\Theta}} + \left(\mathbf{x}_i, \mathbf{x}_j,\mathbf{e}_{j,i}\right) \right), + + where :math:`\square` denotes a differentiable, permutation invariant + function, *e.g.*, sum, mean or max, and :math:`\gamma_{\mathbf{\Theta}}` + and :math:`\phi_{\mathbf{\Theta}}` denote differentiable functions such as + MLPs. + See `here `__ for the accompanying tutorial. + + Args: + aggr (string, optional): The aggregation scheme to use + (:obj:`"add"`, :obj:`"mean"`, :obj:`"max"` or :obj:`None`). + (default: :obj:`"add"`) + flow (string, optional): The flow direction of message passing + (:obj:`"source_to_target"` or :obj:`"target_to_source"`). + (default: :obj:`"source_to_target"`) + node_dim (int, optional): The axis along which to propagate. + (default: :obj:`-2`) + """ + + special_args: Set[str] = { + 'edge_index', 'adj_t', 'edge_index_i', 'edge_index_j', 'size', + 'size_i', 'size_j', 'ptr', 'index', 'dim_size' + } + + def __init__(self, aggr: Optional[str] = "add", + flow: str = "source_to_target", node_dim: int = -2): + + super(MessagePassing, self).__init__() + + self.aggr = aggr + assert self.aggr in ['add', 'mean', 'max', None] + + self.flow = flow + assert self.flow in ['source_to_target', 'target_to_source'] + + self.node_dim = node_dim + + self.inspector = Inspector(self) + self.inspector.inspect(self.message) + self.inspector.inspect(self.aggregate, pop_first=True) + self.inspector.inspect(self.message_and_aggregate, pop_first=True) + self.inspector.inspect(self.update, pop_first=True) + + self.__user_args__ = self.inspector.keys( + ['message', 'aggregate', 'update']).difference(self.special_args) + self.__fused_user_args__ = self.inspector.keys( + ['message_and_aggregate', 'update']).difference(self.special_args) + + # Support for "fused" message passing. + self.fuse = self.inspector.implements('message_and_aggregate') + + # Support for GNNExplainer. + self.__explain__ = False + self.__edge_mask__ = None + + def __check_input__(self, edge_index, size): + the_size: List[Optional[int]] = [None, None] + + if isinstance(edge_index, Tensor): + assert edge_index.dtype == torch.long + assert edge_index.dim() == 2 + assert edge_index.size(0) == 2 + if size is not None: + the_size[0] = size[0] + the_size[1] = size[1] + return the_size + + elif isinstance(edge_index, SparseTensor): + if self.flow == 'target_to_source': + raise ValueError( + ('Flow direction "target_to_source" is invalid for ' + 'message propagation via `torch_sparse.SparseTensor`. If ' + 'you really want to make use of a reverse message ' + 'passing flow, pass in the transposed sparse tensor to ' + 'the message passing module, e.g., `adj_t.t()`.')) + the_size[0] = edge_index.sparse_size(1) + the_size[1] = edge_index.sparse_size(0) + return the_size + + raise ValueError( + ('`MessagePassing.propagate` only supports `torch.LongTensor` of ' + 'shape `[2, num_messages]` or `torch_sparse.SparseTensor` for ' + 'argument `edge_index`.')) + + def __set_size__(self, size: List[Optional[int]], dim: int, src: Tensor): + the_size = size[dim] + if the_size is None: + size[dim] = src.size(self.node_dim) + elif the_size != src.size(self.node_dim): + raise ValueError( + (f'Encountered tensor with size {src.size(self.node_dim)} in ' + f'dimension {self.node_dim}, but expected size {the_size}.')) + + def __lift__(self, src, edge_index, dim): + if isinstance(edge_index, Tensor): + index = edge_index[dim] + return src.index_select(self.node_dim, index) + elif isinstance(edge_index, SparseTensor): + if dim == 1: + rowptr = edge_index.storage.rowptr() + rowptr = expand_left(rowptr, dim=self.node_dim, dims=src.dim()) + return gather_csr(src, rowptr) + elif dim == 0: + col = edge_index.storage.col() + return src.index_select(self.node_dim, col) + raise ValueError + + def __collect__(self, args, edge_index, size, kwargs): + i, j = (1, 0) if self.flow == 'source_to_target' else (0, 1) + + out = {} + for arg in args: + if arg[-2:] not in ['_i', '_j']: + out[arg] = kwargs.get(arg, Parameter.empty) + else: + dim = 0 if arg[-2:] == '_j' else 1 + data = kwargs.get(arg[:-2], Parameter.empty) + + if isinstance(data, (tuple, list)): + assert len(data) == 2 + if isinstance(data[1 - dim], Tensor): + self.__set_size__(size, 1 - dim, data[1 - dim]) + data = data[dim] + + if isinstance(data, Tensor): + self.__set_size__(size, dim, data) + data = self.__lift__(data, edge_index, + j if arg[-2:] == '_j' else i) + + out[arg] = data + + if isinstance(edge_index, Tensor): + out['adj_t'] = None + out['edge_index'] = edge_index + out['edge_index_i'] = edge_index[i] + out['edge_index_j'] = edge_index[j] + out['ptr'] = None + elif isinstance(edge_index, SparseTensor): + out['adj_t'] = edge_index + out['edge_index'] = None + out['edge_index_i'] = edge_index.storage.row() + out['edge_index_j'] = edge_index.storage.col() + out['ptr'] = edge_index.storage.rowptr() + out['edge_weight'] = edge_index.storage.value() + out['edge_attr'] = edge_index.storage.value() + out['edge_type'] = edge_index.storage.value() + + out['index'] = out['edge_index_i'] + out['size'] = size + out['size_i'] = size[1] or size[0] + out['size_j'] = size[0] or size[1] + out['dim_size'] = out['size_i'] + + return out + + def propagate(self, edge_index: Adj, size: Size = None, **kwargs): + r"""The initial call to start propagating messages. + + Args: + edge_index (Tensor or SparseTensor): A :obj:`torch.LongTensor` or a + :obj:`torch_sparse.SparseTensor` that defines the underlying + graph connectivity/message passing flow. + :obj:`edge_index` holds the indices of a general (sparse) + assignment matrix of shape :obj:`[N, M]`. + If :obj:`edge_index` is of type :obj:`torch.LongTensor`, its + shape must be defined as :obj:`[2, num_messages]`, where + messages from nodes in :obj:`edge_index[0]` are sent to + nodes in :obj:`edge_index[1]` + (in case :obj:`flow="source_to_target"`). + If :obj:`edge_index` is of type + :obj:`torch_sparse.SparseTensor`, its sparse indices + :obj:`(row, col)` should relate to :obj:`row = edge_index[1]` + and :obj:`col = edge_index[0]`. + The major difference between both formats is that we need to + input the *transposed* sparse adjacency matrix into + :func:`propagate`. + size (tuple, optional): The size :obj:`(N, M)` of the assignment + matrix in case :obj:`edge_index` is a :obj:`LongTensor`. + If set to :obj:`None`, the size will be automatically inferred + and assumed to be quadratic. + This argument is ignored in case :obj:`edge_index` is a + :obj:`torch_sparse.SparseTensor`. (default: :obj:`None`) + **kwargs: Any additional data which is needed to construct and + aggregate messages, and to update node embeddings. + """ + size = self.__check_input__(edge_index, size) + + # Run "fused" message and aggregation (if applicable). + if (isinstance(edge_index, SparseTensor) and self.fuse + and not self.__explain__): + coll_dict = self.__collect__(self.__fused_user_args__, edge_index, + size, kwargs) + + msg_aggr_kwargs = self.inspector.distribute( + 'message_and_aggregate', coll_dict) + out = self.message_and_aggregate(edge_index, **msg_aggr_kwargs) + + update_kwargs = self.inspector.distribute('update', coll_dict) + return self.update(out, **update_kwargs) + + # Otherwise, run both functions in separation. + elif isinstance(edge_index, Tensor) or not self.fuse: + coll_dict = self.__collect__(self.__user_args__, edge_index, size, + kwargs) + + msg_kwargs = self.inspector.distribute('message', coll_dict) + out = self.message(**msg_kwargs) + + # For `GNNExplainer`, we require a separate message and aggregate + # procedure since this allows us to inject the `edge_mask` into the + # message passing computation scheme. + if self.__explain__: + edge_mask = self.__edge_mask__.sigmoid() + # Some ops add self-loops to `edge_index`. We need to do the + # same for `edge_mask` (but do not train those). + if out.size(self.node_dim) != edge_mask.size(0): + loop = edge_mask.new_ones(size[0]) + edge_mask = torch.cat([edge_mask, loop], dim=0) + assert out.size(self.node_dim) == edge_mask.size(0) + out = out * edge_mask.view([-1] + [1] * (out.dim() - 1)) + + aggr_kwargs = self.inspector.distribute('aggregate', coll_dict) + out = self.aggregate(out, **aggr_kwargs) + + update_kwargs = self.inspector.distribute('update', coll_dict) + return self.update(out, **update_kwargs) + + def message(self, x_j: Tensor) -> Tensor: + r"""Constructs messages from node :math:`j` to node :math:`i` + in analogy to :math:`\phi_{\mathbf{\Theta}}` for each edge in + :obj:`edge_index`. + This function can take any argument as input which was initially + passed to :meth:`propagate`. + Furthermore, tensors passed to :meth:`propagate` can be mapped to the + respective nodes :math:`i` and :math:`j` by appending :obj:`_i` or + :obj:`_j` to the variable name, *.e.g.* :obj:`x_i` and :obj:`x_j`. + """ + return x_j + + def aggregate(self, inputs: Tensor, index: Tensor, + ptr: Optional[Tensor] = None, + dim_size: Optional[int] = None) -> Tensor: + r"""Aggregates messages from neighbors as + :math:`\square_{j \in \mathcal{N}(i)}`. + + Takes in the output of message computation as first argument and any + argument which was initially passed to :meth:`propagate`. + + By default, this function will delegate its call to scatter functions + that support "add", "mean" and "max" operations as specified in + :meth:`__init__` by the :obj:`aggr` argument. + """ + if ptr is not None: + ptr = expand_left(ptr, dim=self.node_dim, dims=inputs.dim()) + return segment_csr(inputs, ptr, reduce=self.aggr) + else: + return scatter(inputs, index, dim=self.node_dim, dim_size=dim_size, + reduce=self.aggr) + + def message_and_aggregate(self, adj_t: SparseTensor) -> Tensor: + r"""Fuses computations of :func:`message` and :func:`aggregate` into a + single function. + If applicable, this saves both time and memory since messages do not + explicitly need to be materialized. + This function will only gets called in case it is implemented and + propagation takes place based on a :obj:`torch_sparse.SparseTensor`. + """ + raise NotImplementedError + + def update(self, inputs: Tensor) -> Tensor: + r"""Updates node embeddings in analogy to + :math:`\gamma_{\mathbf{\Theta}}` for each node + :math:`i \in \mathcal{V}`. + Takes in the output of aggregation as first argument and any argument + which was initially passed to :meth:`propagate`. + """ + return inputs + + @torch.jit.unused + def jittable(self, typing: Optional[str] = None): + r"""Analyzes the :class:`MessagePassing` instance and produces a new + jittable module. + + Args: + typing (string, optional): If given, will generate a concrete + instance with :meth:`forward` types based on :obj:`typing`, + *e.g.*: :obj:`"(Tensor, Optional[Tensor]) -> Tensor"`. + """ + # Find and parse `propagate()` types to format `{arg1: type1, ...}`. + if hasattr(self, 'propagate_type'): + prop_types = { + k: sanitize(str(v)) + for k, v in self.propagate_type.items() + } + else: + source = inspect.getsource(self.__class__) + match = re.search(r'#\s*propagate_type:\s*\((.*)\)', source) + if match is None: + raise TypeError( + 'TorchScript support requires the definition of the types ' + 'passed to `propagate()`. Please specificy them via\n\n' + 'propagate_type = {"arg1": type1, "arg2": type2, ... }\n\n' + 'or via\n\n' + '# propagate_type: (arg1: type1, arg2: type2, ...)\n\n' + 'inside the `MessagePassing` module.') + prop_types = split_types_repr(match.group(1)) + prop_types = dict([re.split(r'\s*:\s*', t) for t in prop_types]) + + # Parse `__collect__()` types to format `{arg:1, type1, ...}`. + collect_types = self.inspector.types( + ['message', 'aggregate', 'update']) + + # Collect `forward()` header, body and @overload types. + forward_types = parse_types(self.forward) + forward_types = [resolve_types(*types) for types in forward_types] + forward_types = list(chain.from_iterable(forward_types)) + + keep_annotation = len(forward_types) < 2 + forward_header = func_header_repr(self.forward, keep_annotation) + forward_body = func_body_repr(self.forward, keep_annotation) + + if keep_annotation: + forward_types = [] + elif typing is not None: + forward_types = [] + forward_body = 8 * ' ' + f'# type: {typing}\n{forward_body}' + + root = os.path.dirname(osp.realpath(__file__)) + with open(osp.join(root, 'message_passing.jinja'), 'r') as f: + template = Template(f.read()) + + uid = uuid1().hex[:6] + cls_name = f'{self.__class__.__name__}Jittable_{uid}' + jit_module_repr = template.render( + uid=uid, + module=str(self.__class__.__module__), + cls_name=cls_name, + parent_cls_name=self.__class__.__name__, + prop_types=prop_types, + collect_types=collect_types, + user_args=self.__user_args__, + forward_header=forward_header, + forward_types=forward_types, + forward_body=forward_body, + msg_args=self.inspector.keys(['message']), + aggr_args=self.inspector.keys(['aggregate']), + msg_and_aggr_args=self.inspector.keys(['message_and_aggregate']), + update_args=self.inspector.keys(['update']), + check_input=inspect.getsource(self.__check_input__)[:-1], + lift=inspect.getsource(self.__lift__)[:-1], + ) + + # Instantiate a class from the rendered JIT module representation. + cls = class_from_module_repr(cls_name, jit_module_repr) + module = cls.__new__(cls) + module.__dict__ = self.__dict__.copy() + module.jittable = None + + return module diff --git a/autogl/module/nas/space/gasso_space/sage_conv.py b/autogl/module/nas/space/gasso_space/sage_conv.py new file mode 100644 index 0000000..0edf1d6 --- /dev/null +++ b/autogl/module/nas/space/gasso_space/sage_conv.py @@ -0,0 +1,92 @@ +from typing import Union, Tuple +from torch_geometric.typing import (OptPairTensor, Adj, Size, NoneType, + OptTensor) + +from torch import Tensor +from torch.nn import Linear +import torch.nn.functional as F +from torch_sparse import SparseTensor, matmul +from torch_geometric.nn.conv import MessagePassing + + +class SAGEConv(MessagePassing): + r"""The GraphSAGE operator from the `"Inductive Representation Learning on + Large Graphs" `_ paper + + .. math:: + \mathbf{x}^{\prime}_i = \mathbf{W}_1 \mathbf{x}_i + \mathbf{W_2} \cdot + \mathrm{mean}_{j \in \mathcal{N(i)}} \mathbf{x}_j + + Args: + in_channels (int or tuple): Size of each input sample. A tuple + corresponds to the sizes of source and target dimensionalities. + out_channels (int): Size of each output sample. + normalize (bool, optional): If set to :obj:`True`, output features + will be :math:`\ell_2`-normalized, *i.e.*, + :math:`\frac{\mathbf{x}^{\prime}_i} + {\| \mathbf{x}^{\prime}_i \|_2}`. + (default: :obj:`False`) + bias (bool, optional): If set to :obj:`False`, the layer will not learn + an additive bias. (default: :obj:`True`) + **kwargs (optional): Additional arguments of + :class:`torch_geometric.nn.conv.MessagePassing`. + """ + def __init__(self, in_channels: Union[int, Tuple[int, int]], + out_channels: int, normalize: bool = False, + bias: bool = True, **kwargs): # yapf: disable + kwargs.setdefault('aggr', 'mean') + super(SAGEConv, self).__init__(**kwargs) + + self.in_channels = in_channels + self.out_channels = out_channels + self.normalize = normalize + + if isinstance(in_channels, int): + in_channels = (in_channels, in_channels) + + self.lin_l = Linear(in_channels[0], out_channels, bias=bias) + self.lin_r = Linear(in_channels[1], out_channels, bias=False) + + self.reset_parameters() + + def reset_parameters(self): + self.lin_l.reset_parameters() + self.lin_r.reset_parameters() + + def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj, edge_weight: OptTensor = None, + size: Size = None) -> Tensor: + """""" + if isinstance(x, Tensor): + x: OptPairTensor = (x, x) + + # propagate_type: (x: OptPairTensor) + out = self.propagate(edge_index, x=x, edge_weight = edge_weight, size=size) + out = self.lin_l(out) + + x_r = x[1] + if x_r is not None: + out += self.lin_r(x_r) + + if self.normalize: + out = F.normalize(out, p=2., dim=-1) + + return out + + #def message(self, x_j: Tensor) -> Tensor: + # return x_j + def message(self, x_j: Tensor, edge_weight: OptTensor) -> Tensor: + if edge_weight is None: + return x_j + else: + return edge_weight.view(-1, 1) * x_j + + + + def message_and_aggregate(self, adj_t: SparseTensor, + x: OptPairTensor) -> Tensor: + adj_t = adj_t.set_value(None, layout=None) + return matmul(adj_t, x[0], reduce=self.aggr) + + def __repr__(self): + return '{}({}, {})'.format(self.__class__.__name__, self.in_channels, + self.out_channels) diff --git a/autogl/module/nas/space/gasso_space/utils/__init__.py b/autogl/module/nas/space/gasso_space/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/autogl/module/nas/space/gasso_space/utils/helpers.py b/autogl/module/nas/space/gasso_space/utils/helpers.py new file mode 100644 index 0000000..5961a44 --- /dev/null +++ b/autogl/module/nas/space/gasso_space/utils/helpers.py @@ -0,0 +1,7 @@ +import torch + + +def expand_left(src: torch.Tensor, dim: int, dims: int) -> torch.Tensor: + for _ in range(dims + dim if dim < 0 else dim): + src = src.unsqueeze(0) + return src diff --git a/autogl/module/nas/space/gasso_space/utils/inspector.py b/autogl/module/nas/space/gasso_space/utils/inspector.py new file mode 100644 index 0000000..d9a2fae --- /dev/null +++ b/autogl/module/nas/space/gasso_space/utils/inspector.py @@ -0,0 +1,86 @@ +import re +import inspect +from collections import OrderedDict +from typing import Dict, List, Any, Optional, Callable, Set + +from .typing import parse_types + + +class Inspector(object): + def __init__(self, base_class: Any): + self.base_class: Any = base_class + self.params: Dict[str, Dict[str, Any]] = {} + + def inspect(self, func: Callable, + pop_first: bool = False) -> Dict[str, Any]: + params = inspect.signature(func).parameters + params = OrderedDict(params) + if pop_first: + params.popitem(last=False) + self.params[func.__name__] = params + + def keys(self, func_names: Optional[List[str]] = None) -> Set[str]: + keys = [] + for func in func_names or list(self.params.keys()): + keys += self.params[func].keys() + return set(keys) + + def __implements__(self, cls, func_name: str) -> bool: + if cls.__name__ == 'MessagePassing': + return False + if func_name in cls.__dict__.keys(): + return True + return any(self.__implements__(c, func_name) for c in cls.__bases__) + + def implements(self, func_name: str) -> bool: + return self.__implements__(self.base_class.__class__, func_name) + + def types(self, func_names: Optional[List[str]] = None) -> Dict[str, str]: + out: Dict[str, str] = {} + for func_name in func_names or list(self.params.keys()): + func = getattr(self.base_class, func_name) + arg_types = parse_types(func)[0][0] + for key in self.params[func_name].keys(): + if key in out and out[key] != arg_types[key]: + raise ValueError( + (f'Found inconsistent types for argument {key}. ' + f'Expected type {out[key]} but found type ' + f'{arg_types[key]}.')) + out[key] = arg_types[key] + return out + + def distribute(self, func_name, kwargs: Dict[str, Any]): + out = {} + for key, param in self.params[func_name].items(): + data = kwargs.get(key, inspect.Parameter.empty) + if data is inspect.Parameter.empty: + if param.default is inspect.Parameter.empty: + raise TypeError(f'Required parameter {key} is empty.') + data = param.default + out[key] = data + return out + + +def func_header_repr(func: Callable, keep_annotation: bool = True) -> str: + source = inspect.getsource(func) + signature = inspect.signature(func) + + if keep_annotation: + return ''.join(re.split(r'(\).*?:.*?\n)', source, + maxsplit=1)[:2]).strip() + + params_repr = ['self'] + for param in signature.parameters.values(): + params_repr.append(param.name) + if param.default is not inspect.Parameter.empty: + params_repr[-1] += f'={param.default}' + + return f'def {func.__name__}({", ".join(params_repr)}):' + + +def func_body_repr(func: Callable, keep_annotation: bool = True) -> str: + source = inspect.getsource(func) + body_repr = re.split(r'\).*?:.*?\n', source, maxsplit=1)[1] + if not keep_annotation: + body_repr = re.sub(r'\s*# type:.*\n', '', body_repr) + return body_repr diff --git a/autogl/module/nas/space/gasso_space/utils/jit.py b/autogl/module/nas/space/gasso_space/utils/jit.py new file mode 100644 index 0000000..150b348 --- /dev/null +++ b/autogl/module/nas/space/gasso_space/utils/jit.py @@ -0,0 +1,19 @@ +import sys +import os.path as osp +from getpass import getuser +from tempfile import NamedTemporaryFile as TempFile, gettempdir +from importlib.util import module_from_spec, spec_from_file_location + +from torch_geometric.data.makedirs import makedirs + + +def class_from_module_repr(cls_name, module_repr): + path = osp.join(gettempdir(), f'{getuser()}_pyg_jit') + makedirs(path) + with TempFile(mode='w+', suffix='.py', delete=False, dir=path) as f: + f.write(module_repr) + spec = spec_from_file_location(cls_name, f.name) + mod = module_from_spec(spec) + sys.modules[cls_name] = mod + spec.loader.exec_module(mod) + return getattr(mod, cls_name) diff --git a/autogl/module/nas/space/gasso_space/utils/typing.py b/autogl/module/nas/space/gasso_space/utils/typing.py new file mode 100644 index 0000000..bb5f228 --- /dev/null +++ b/autogl/module/nas/space/gasso_space/utils/typing.py @@ -0,0 +1,107 @@ +import re +import inspect +import pyparsing as pp +from itertools import product +from collections import OrderedDict +from typing import Callable, Tuple, Dict, List + + +def split_types_repr(types_repr: str) -> List[str]: + out = [] + i = depth = 0 + for j, char in enumerate(types_repr): + if char == '[': + depth += 1 + elif char == ']': + depth -= 1 + elif char == ',' and depth == 0: + out.append(types_repr[i:j].strip()) + i = j + 1 + out.append(types_repr[i:].strip()) + return out + + +def sanitize(type_repr: str): + type_repr = re.sub(r'', r'\1', type_repr) + type_repr = type_repr.replace('typing.', '') + type_repr = type_repr.replace('torch_sparse.tensor.', '') + type_repr = type_repr.replace('Adj', 'Union[Tensor, SparseTensor]') + + # Replace `Union[..., NoneType]` by `Optional[...]`. + sexp = pp.nestedExpr(opener='[', closer=']') + tree = sexp.parseString(f'[{type_repr.replace(",", " ")}]').asList()[0] + + def union_to_optional_(tree): + for i in range(len(tree)): + e, n = tree[i], tree[i + 1] if i + 1 < len(tree) else [] + if e == 'Union' and n[-1] == 'NoneType': + tree[i] = 'Optional' + tree[i + 1] = tree[i + 1][:-1] + elif e == 'Union' and 'NoneType' in n: + idx = n.index('NoneType') + n[idx] = [n[idx - 1]] + n[idx - 1] = 'Optional' + elif isinstance(e, list): + tree[i] = union_to_optional_(e) + return tree + + tree = union_to_optional_(tree) + type_repr = re.sub(r'\'|\"', '', str(tree)[1:-1]).replace(', [', '[') + + return type_repr + + +def param_type_repr(param) -> str: + if param.annotation is inspect.Parameter.empty: + return 'torch.Tensor' + return sanitize(re.split(r':|='.strip(), str(param))[1]) + + +def return_type_repr(signature) -> str: + return_type = signature.return_annotation + if return_type is inspect.Parameter.empty: + return 'torch.Tensor' + elif str(return_type)[:6] != ' List[Tuple[Dict[str, str], str]]: + source = inspect.getsource(func) + signature = inspect.signature(func) + + # Parse `# type: (...) -> ...` annotation. Note that it is allowed to pass + # multiple `# type:` annotations in `forward()`. + iterator = re.finditer(r'#\s*type:\s*\((.*)\)\s*->\s*(.*)\s*\n', source) + matches = list(iterator) + + if len(matches) > 0: + out = [] + args = list(signature.parameters.keys()) + for match in matches: + arg_types_repr, return_type = match.groups() + arg_types = split_types_repr(arg_types_repr) + arg_types = OrderedDict((k, v) for k, v in zip(args, arg_types)) + return_type = return_type.split('#')[0].strip() + out.append((arg_types, return_type)) + return out + + # Alternatively, parse annotations using the inspected signature. + else: + ps = signature.parameters + arg_types = OrderedDict((k, param_type_repr(v)) for k, v in ps.items()) + return [(arg_types, return_type_repr(signature))] + + +def resolve_types(arg_types: Dict[str, str], + return_type_repr: str) -> List[Tuple[List[str], str]]: + out = [] + for type_repr in arg_types.values(): + if type_repr[:5] == 'Union': + out.append(split_types_repr(type_repr[6:-1])) + else: + out.append([type_repr]) + return [(x, return_type_repr) for x in product(*out)] diff --git a/configs/nodeclf_nas_gasso.yml b/configs/nodeclf_nas_gasso.yml new file mode 100644 index 0000000..4a5d985 --- /dev/null +++ b/configs/nodeclf_nas_gasso.yml @@ -0,0 +1,40 @@ +ensemble: + name: null +feature: +- name: NormalizeFeatures +hpo: + max_evals: 10 + name: random +nas: + space: + name: gassospace + hidden_dim: 64 + layer_number: 2 + algorithm: + name: gasso + num_epochs: 250 + estimator: + name: oneshot +models: [] +trainer: + hp_space: + - maxValue: 300 + minValue: 100 + parameterName: max_epoch + scalingType: LINEAR + type: INTEGER + - maxValue: 30 + minValue: 10 + parameterName: early_stopping_round + scalingType: LINEAR + type: INTEGER + - maxValue: 0.05 + minValue: 0.01 + parameterName: lr + scalingType: LOG + type: DOUBLE + - maxValue: 0.0005 + minValue: 5.0e-05 + parameterName: weight_decay + scalingType: LOG + type: DOUBLE diff --git a/examples/gasso_test.py b/examples/gasso_test.py new file mode 100644 index 0000000..2e5166b --- /dev/null +++ b/examples/gasso_test.py @@ -0,0 +1,24 @@ +import os +os.environ["AUTOGL_BACKEND"] = "pyg" +import sys +sys.path.append('../') +from autogl.datasets import build_dataset_from_name +from autogl.solver import AutoNodeClassifier +from autogl.module.train import Acc +from autogl.solver.utils import set_seed +import argparse + +if __name__ == '__main__': + set_seed(202106) + parser = argparse.ArgumentParser() + parser.add_argument('--config', type=str, default='../configs/nodeclf_nas_gasso.yml') + parser.add_argument('--dataset', choices=['cora', 'citeseer', 'pubmed'], default='cora', type=str) + + args = parser.parse_args() + + dataset = build_dataset_from_name('cora', path = "~/AGL/") + solver = AutoNodeClassifier.from_config(args.config) + solver.fit(dataset) + solver.get_leaderboard().show() + out = solver.predict_proba() + print('acc on dataset', Acc.evaluate(out, dataset[0].y[dataset[0].test_mask].detach().numpy())) diff --git a/test/nas/node_classification.py b/test/nas/node_classification.py index e523e3d..33c3acc 100644 --- a/test/nas/node_classification.py +++ b/test/nas/node_classification.py @@ -81,16 +81,16 @@ if __name__ == "__main__": di = bk_feat(data).shape[1] do = len(np.unique(bk_label(data))) - print("Random search + graphnas ") - space = GraphNasNodeClassificationSpace().cuda() + print("Random search + singlepath ") + space = SinglePathNodeClassificationSpace().cuda() space.instantiate(input_dim=di, output_dim=do) esti = OneShotEstimator() algo = RandomSearch(num_epochs=10) model = algo.search(space, dataset, esti) test_model(model, data, True) - print("Random search + singlepath ") - space = SinglePathNodeClassificationSpace().cuda() + print("Random search + graphnas ") + space = GraphNasNodeClassificationSpace().cuda() space.instantiate(input_dim=di, output_dim=do) esti = OneShotEstimator() algo = RandomSearch(num_epochs=10) @@ -113,14 +113,6 @@ if __name__ == "__main__": model = algo.search(space, dataset, esti) test_model(model, data, True) - print("enas + graphnas ") - space = GraphNasNodeClassificationSpace().cuda() - space.instantiate(input_dim=di, output_dim=do) - esti = OneShotEstimator() - algo = Enas(num_epochs=10) - model = algo.search(space, dataset, esti) - test_model(model, data, True) - print("darts + graphnas ") space = GraphNasNodeClassificationSpace(con_ops=['concat']).cuda() space.instantiate(input_dim=di, output_dim=do) From 74a22ad81c3c4262352ba8f64bfe9050bf815539 Mon Sep 17 00:00:00 2001 From: generall Date: Mon, 6 Dec 2021 16:00:04 +0800 Subject: [PATCH 03/14] improve gasso --- autogl/module/nas/algorithm/gasso.py | 9 ++++++--- autogl/module/nas/estimator/one_shot.py | 5 +++-- autogl/module/nas/space/gasso.py | 5 +++-- examples/gasso_test.py | 5 +++-- 4 files changed, 15 insertions(+), 9 deletions(-) diff --git a/autogl/module/nas/algorithm/gasso.py b/autogl/module/nas/algorithm/gasso.py index 28ec33f..42345cc 100644 --- a/autogl/module/nas/algorithm/gasso.py +++ b/autogl/module/nas/algorithm/gasso.py @@ -59,7 +59,6 @@ class Gasso(BaseNAS): diff = (logits[e1] - logits[e2]).pow(2).sum(1) smooth = (diff * torch.sigmoid(ew)).sum() dist = (ew * ew).sum() - #print(smooth,dist) loss += self.lamb * smooth + dist optimizer.zero_grad() @@ -100,8 +99,9 @@ class Gasso(BaseNAS): t_total = time.time() for epoch in range(self.num_epochs): + self.space.train() self.optimizer.zero_grad() - acc1, loss = self._infer(self.space, data, self.estimator, "train") + _, loss = self._infer(self.space, data, self.estimator, "train") loss.backward() self.optimizer.step() @@ -110,15 +110,18 @@ class Gasso(BaseNAS): self.train_stru(self.space, self.stru_optimizer, data) self.arch_optimizer.zero_grad() - acc2, loss = self._infer(self.space, data, self.estimator, "train") + _, loss = self._infer(self.space, data, self.estimator, "train") loss.backward() self.arch_optimizer.step() + self.space.eval() + train_acc, _ = self._infer(self.space, data, self.estimator, "train") val_acc, val_loss = self._infer(self.space, data, self.estimator, "val") if val_loss < min_val_loss: min_val_loss = val_loss best_performance = val_acc self.space.keep_prediction() + #print("acc:" + str(train_acc) + " val_acc" + str(val_acc)) return best_performance, min_val_loss diff --git a/autogl/module/nas/estimator/one_shot.py b/autogl/module/nas/estimator/one_shot.py index e598959..0081ee3 100644 --- a/autogl/module/nas/estimator/one_shot.py +++ b/autogl/module/nas/estimator/one_shot.py @@ -1,4 +1,5 @@ import torch.nn.functional as F +import torch from . import register_nas_estimator from ..space import BaseSpace @@ -7,7 +8,6 @@ from ..backend import * from ...train.evaluation import Acc from ..utils import get_hardware_aware_metric - @register_nas_estimator("oneshot") class OneShotEstimator(BaseEstimator): """ @@ -30,7 +30,8 @@ class OneShotEstimator(BaseEstimator): y = label[mask] loss = getattr(F, self.loss_f)(pred, y) - probs = F.softmax(pred, dim=1).detach().cpu().numpy() + #probs = F.softmax(pred, dim=1).detach().cpu().numpy() + probs = pred.detach().cpu().numpy() y = y.cpu() metrics = [eva.evaluate(probs, y) for eva in self.evaluation] return metrics, loss diff --git a/autogl/module/nas/space/gasso.py b/autogl/module/nas/space/gasso.py index 5a16fe7..30ab8cf 100644 --- a/autogl/module/nas/space/gasso.py +++ b/autogl/module/nas/space/gasso.py @@ -191,7 +191,7 @@ class GassoSpace(BaseSpace): self, hidden_dim: _typ.Optional[int] = 64, layer_number: _typ.Optional[int] = 2, - dropout: _typ.Optional[float] = 0.2, + dropout: _typ.Optional[float] = 0.8, input_dim: _typ.Optional[int] = None, output_dim: _typ.Optional[int] = None, ops: _typ.Tuple = gnn_list, @@ -210,7 +210,7 @@ class GassoSpace(BaseSpace): self, hidden_dim: _typ.Optional[int] = 64, layer_number: _typ.Optional[int] = 2, - dropout: _typ.Optional[float] = 0.2, + dropout: _typ.Optional[float] = 0.8, input_dim: _typ.Optional[int] = None, output_dim: _typ.Optional[int] = None, ops: _typ.Tuple = gnn_list, @@ -233,6 +233,7 @@ class GassoSpace(BaseSpace): self.initialize_alphas() + #def forward(self, x, adjs): def forward(self, data): if self.use_forward: x, adjs = data.x, data.adj diff --git a/examples/gasso_test.py b/examples/gasso_test.py index 2e5166b..39efea6 100644 --- a/examples/gasso_test.py +++ b/examples/gasso_test.py @@ -12,11 +12,12 @@ if __name__ == '__main__': set_seed(202106) parser = argparse.ArgumentParser() parser.add_argument('--config', type=str, default='../configs/nodeclf_nas_gasso.yml') - parser.add_argument('--dataset', choices=['cora', 'citeseer', 'pubmed'], default='cora', type=str) + parser.add_argument('--dataset', choices=['cora', 'citeseer', 'pubmed'], default='citeseer', type=str) args = parser.parse_args() - dataset = build_dataset_from_name('cora', path = "~/AGL/") + #dataset = build_dataset_from_name(args.dataset, path = "/DATA/DATANAS1/qinyj/enhgnas/") + dataset = build_dataset_from_name(args.dataset, path = "~/AGL/") solver = AutoNodeClassifier.from_config(args.config) solver.fit(dataset) solver.get_leaderboard().show() From 3e0efbcf5f80f3ab86a597d4f8dec680c583015d Mon Sep 17 00:00:00 2001 From: generall Date: Mon, 6 Dec 2021 16:01:33 +0800 Subject: [PATCH 04/14] reduce gasso --- autogl/module/nas/space/gasso.py | 34 -------------------------------- 1 file changed, 34 deletions(-) diff --git a/autogl/module/nas/space/gasso.py b/autogl/module/nas/space/gasso.py index 30ab8cf..0397cb5 100644 --- a/autogl/module/nas/space/gasso.py +++ b/autogl/module/nas/space/gasso.py @@ -270,40 +270,6 @@ class GassoSpace(BaseSpace): def arch_parameters(self): return self.alphas_normal - def parse_gene_force(self, alphas): - gene = [] - n = 2 - max_num_edges = 1 - start = 0 - mat = F.softmax(torch.stack(alphas, dim=0), dim=-1).detach() - importance = torch.sum(mat[:, 1:], dim=-1).cpu() - for i in range(self.steps): - #end = start + 2 - end = start + n - num_edges_to_select = max_num_edges - if num_edges_to_select > 0: - post_select_edges = torch.topk(importance[start: end], k=num_edges_to_select).indices + start - else: - post_select_edges = [] - for j in range(start, end): - if num_edges_to_select <= 0: - raise Exception("Unknown errors") - else: - if j in post_select_edges: - idx = torch.argmax(alphas[j][1:]) + 1 - gene.append((self.ops[idx], j - start)) - start = end - n += 1 - - return gene - - def get_genotype(self, force=True): - gene_normal = self.parse_gene_force(self.alphas_normal) - n = 2 - concat = range(n, self.steps + n) - genotype = Genotype(normal=gene_normal, normal_concat=concat) - return genotype - def parse_model(self, selection, device) -> BaseModel: self.use_forward = False return self.wrap(device) From 3c3dc06377bf349c8af04e97011c3899c93cebce Mon Sep 17 00:00:00 2001 From: wondergo2017 Date: Fri, 17 Dec 2021 18:59:26 +0800 Subject: [PATCH 05/14] add autoattend --- autogl/module/nas/space/autoattend.py | 204 ++++++++++++++++++ .../nas/space/autoattend_space/__init__.py | 0 .../nas/space/autoattend_space/operations.py | 201 +++++++++++++++++ .../module/nas/space/autoattend_space/ops1.py | 19 ++ .../module/nas/space/autoattend_space/ops2.py | 13 ++ autogl/module/nas/space/graph_nas_macro.py | 6 +- test/nas/node_classification.py | 120 ++++++----- 7 files changed, 506 insertions(+), 57 deletions(-) create mode 100644 autogl/module/nas/space/autoattend.py create mode 100644 autogl/module/nas/space/autoattend_space/__init__.py create mode 100644 autogl/module/nas/space/autoattend_space/operations.py create mode 100644 autogl/module/nas/space/autoattend_space/ops1.py create mode 100644 autogl/module/nas/space/autoattend_space/ops2.py diff --git a/autogl/module/nas/space/autoattend.py b/autogl/module/nas/space/autoattend.py new file mode 100644 index 0000000..a33b375 --- /dev/null +++ b/autogl/module/nas/space/autoattend.py @@ -0,0 +1,204 @@ +# codes in this file are reproduced from https://github.com/GraphNAS/GraphNAS with some changes. +from nni.nas.pytorch.mutables import Mutable +import typing as _typ +import torch + +import torch.nn.functional as F +from nni.nas.pytorch import mutables + +from . import register_nas_space +from .base import BaseSpace +from ...model import BaseModel +from ..utils import count_parameters, measure_latency + +from torch import nn +from .operation import act_map, gnn_map + +from ..backend import * + +from .autoattend_space.ops1 import OPS as OPS1 +from .autoattend_space.ops2 import OPS as OPS2 +from .autoattend_space.operations import agg_map +OPS = [OPS1, OPS2] + +from nni.nas.pytorch.mutables import Mutable +class MultiLayerChoice(Mutable): + def __init__(self, choices, layer, key): + super(MultiLayerChoice, self).__init__(key) + self.order = layer + self.choices = choices + + def forward(self, *args, **kwargs): + outs = [] + for i in range(len(self.choices)): + out = self.choices[i](*args, **kwargs) + outs.append(out) + outs = torch.stack(outs, dim=0) + return outs + + +@register_nas_space("autoattend") +class AutoAttendNodeClassificationSpace(BaseSpace): + def __init__( + self, + hidden_dim: _typ.Optional[int] = 64, + layer_number: _typ.Optional[int] = 2, + dropout: _typ.Optional[float] = 0.9, + input_dim: _typ.Optional[int] = None, + output_dim: _typ.Optional[int] = None, + ops_type=0, + gnn_ops: _typ.Sequence[_typ.Union[str, _typ.Any] + ] = None, + act_op="tanh", + head=8, + agg_ops=['add','attn'] + ): + super().__init__() + self.layer_number = layer_number + self.hidden_dim = hidden_dim + self.input_dim = input_dim + self.output_dim = output_dim + self.gnn_ops = gnn_ops + self.dropout = dropout + self.act_op = act_op + self.ops_type = ops_type + self.head = head + self.agg_ops = agg_ops + + def instantiate( + self, + hidden_dim: _typ.Optional[int] = None, + layer_number: _typ.Optional[int] = None, + dropout: _typ.Optional[float] = None, + input_dim: _typ.Optional[int] = None, + output_dim: _typ.Optional[int] = None, + ops_type=None, + gnn_ops: _typ.Sequence[_typ.Union[str, _typ.Any]] = None, + act_op=None, + head=None, + agg_ops=None, + # con_ops: _typ.Sequence[_typ.Union[str, _typ.Any]] = None, + ): + super().instantiate() + self.dropout = dropout or self.dropout + self.hidden_dim = hidden_dim or self.hidden_dim + self.layer_number = layer_number or self.layer_number + self.input_dim = input_dim or self.input_dim + self.output_dim = output_dim or self.output_dim + self.gnn_ops = gnn_ops or self.gnn_ops + self.act_op = act_op or self.act_op + self.act = act_map(self.act_op) + self.head = head or self.head + self.ops_type = ops_type or self.ops_type + self.agg_ops=agg_ops or self.agg_ops + PRIMITIVES = list(OPS[self.ops_type].keys()) + self.gnn_map = lambda x, * \ + args, **kwargs: OPS[self.ops_type][x](*args, **kwargs) + self.gnn_ops = self.gnn_ops or PRIMITIVES + self.agg_map = lambda x, * \ + args, **kwargs: agg_map[x](*args, **kwargs) + # self.con_ops = con_ops or self.con_ops + self.preproc0 = nn.Linear(self.input_dim, self.hidden_dim) + # self.preproc1 = nn.Linear(self.input_dim, self.hidden_dim) + node_labels = [] + + # stem path + for layer in range(1, self.layer_number+1): + key = f"stem_{layer}" + self._set_layer_choice(layer, key) + + # side path + for layer in range(1, self.layer_number+1): + choices = [] + key = f"side_{layer}" + for i in range(2): + sub_key = f"{key}_{i}" + choice = self._set_layer_choice(layer, sub_key) + choices.append(choice) + setattr(self, key, MultiLayerChoice(choices, layer, key)) + node_labels.append(key) + + # input + key = f"in_{layer}" + self._set_input_choice(key, + layer, choose_from=node_labels, n_chosen=1, return_mask=False) + + # agg + for layer in range(1, self.layer_number + 1): + key = f"agg_{layer}" + self._set_agg_choice(layer, key=key) + + self._initialized = True + + self.classifier2 = nn.Linear(self.hidden_dim, self.output_dim) + + print(self) + def _set_agg_choice(self,layer,key): + ops=[self.agg_map(op, self.hidden_dim,self.head,self.dropout)for op in self.agg_ops] + choice = self.setLayerChoice( + layer, + ops, + key=key, + ) + setattr(self, key, choice) + return choice + def _set_layer_choice(self, layer, key): + if self.ops_type==0: + ops=[self.gnn_map(op, self.hidden_dim, self.hidden_dim, self.dropout)for op in self.gnn_ops] + elif self.ops_type==1: + ops=[self.gnn_map(op, self.hidden_dim, self.hidden_dim, self.head, self.dropout)for op in self.gnn_ops] + choice = self.setLayerChoice( + layer, + ops, + key=key, + ) + setattr(self, key, choice) + return choice + + def _set_input_choice(self, key, layer,**kwargs): + setattr(self, + key, + self.setInputChoice( + layer, + key=key, + **kwargs + )) + + def forward(self, data): + x = bk_feat(data) + x = F.dropout(x, p=self.dropout, training=self.training) + prev_ = self.preproc0(x) + + side_outs = [] + stem_outs = [] + input = prev_ + for layer in range(1, self.layer_number + 1): + # do layer choice for stem + op = getattr(self, f"stem_{layer}") + print(op) + stem_out = bk_gconv(op, data, input) + stem_out = self.act(stem_out) + # do double layer choice for sides + op = getattr(self, f'side_{layer}') + side_out = bk_gconv(op, data, input) + side_out = self.act(side_out) + + stem_outs.append(stem_out) + side_outs.append(side_out) + + # select input [x1,x2,x3] from side1,side2,stem + side_selected = getattr(self, f"in_{layer}")(side_outs) + input = [stem_outs[-1], side_selected] + # do agg in [add , attn] + agg = getattr(self, f"agg_{layer}") + print(layer,input) + input = bk_gconv(agg,data,input) + + x = self.classifier2(input) + return F.log_softmax(x, dim=1) + + def parse_model(self, selection, device) -> BaseModel: + for i in range(1, self.layer_number + 1): + selection[f'side_{i}'] = None + # return AutoGCN(self.input_dim, self.output_dim, device) + return self.wrap(device).fix(selection) diff --git a/autogl/module/nas/space/autoattend_space/__init__.py b/autogl/module/nas/space/autoattend_space/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/autogl/module/nas/space/autoattend_space/operations.py b/autogl/module/nas/space/autoattend_space/operations.py new file mode 100644 index 0000000..74b7f7f --- /dev/null +++ b/autogl/module/nas/space/autoattend_space/operations.py @@ -0,0 +1,201 @@ + +from torch_geometric.nn import MessagePassing +from torch_geometric.utils import softmax +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch_geometric.nn import GCNConv, SAGEConv, GATConv, ARMAConv, ChebConv, GatedGraphConv, SGConv + +from typing import Union, Tuple, Optional +from torch_geometric.typing import (OptPairTensor, Adj, Size, NoneType, + OptTensor) + +import torch +from torch import Tensor +import torch.nn.functional as F +from torch.nn import Parameter, Linear +from torch_sparse import SparseTensor, set_diag +from torch_geometric.nn.conv import MessagePassing +from torch_geometric.utils import remove_self_loops, add_self_loops, softmax + +from torch_geometric.nn.inits import glorot, zeros +import torch +import torch.nn as nn +import torch.nn.functional as F + +import torch.nn.functional as F +from torch.nn import Parameter +from torch_geometric.nn.inits import glorot, zeros +from torch_geometric.utils import softmax +from torch_scatter import scatter_add +import numpy as np +from ..graph_nas_macro import GeoLayer + +class AggAdd(nn.Module): + def __init__(self, dim, att_head, dropout=0,norm=False,skip_connect=False,*args,**kwargs): + super(AggAdd, self).__init__() + self.dropout = dropout + self.ln_add = nn.BatchNorm1d(dim, track_running_stats=True, affine=True) + self.norm=norm + self.skip_connect=skip_connect + + def forward(self,x,edge_index,*args,**kwargs): + # x=[x_stem,[x_sides]] + norm=self.norm + x1,x2,x3=x[0],x[1][0],x[1][1] + if norm: + return self.ln_add(x1 + x2) + else: + return x1 + x2 + +class AggAttn(MessagePassing) : + def __init__(self, dim, att_head,dropout=0,norm=False,skip_connect=False,*args,**kwargs): + super(AggAttn, self).__init__() + self.dropout = dropout + self.att_head=att_head + self.ln_attn = nn.BatchNorm1d(dim, track_running_stats=True, affine=True) + self.norm=norm + self.skip_connect=skip_connect + def __repr__(self) -> str: + return 'AggAttn(att_head={}, dropout={})'.format(self.att_head, self.dropout) + def forward(self,x,edge_index,*args,**kwargs): + # x=[x_stem,[x_sides]] + # use dot-product attn + x1,x2,x3=x[0],x[1][0],x[1][1] # q,k,v + skip_connect,norm=self.skip_connect,self.norm + + if not skip_connect and not norm: + return self.propagate(edge_index, x1=x1, x2=x2, x3=x3).view(x1.size(0), -1) + + x = self.propagate(edge_index, x1=x1, x2=x2, x3=x3).view(x1.size(0), -1) + if not norm: + return x + + if not skip_connect: + return self.ln_attn(x) + return self.ln_attn(x + x1) + + def message(self, x2_j, x1_i, x3_j, index, ptr): + # x1: query, x2: key, x3: value + node, dim = x1_i.size() + dim_att = dim // self.att_head + x2_j = x2_j.view(node, self.att_head, dim_att) + x1_i = x1_i.view(node, self.att_head, dim_att) + attn = (x2_j * x1_i).sum(dim=-1) / np.sqrt(dim_att) + attn = softmax(attn, index, ptr) + attn = F.dropout(attn, p=self.dropout, training=self.training) + return x3_j.view(node, self.att_head, dim_att) * attn.unsqueeze(-1) + +class GATConv2(MessagePassing): + _alpha: OptTensor + + def __init__(self, in_channels: Union[int, Tuple[int, int]], + out_channels: int, heads: int = 1, concat: bool = False, + negative_slope: float = 0.2, dropout: float = 0., + add_self_loops: bool = True, bias: bool = True, **kwargs): + super(GATConv, self).__init__(aggr='add', node_dim=0, **kwargs) + + self.in_channels = in_channels + self.out_channels = out_channels + self.heads = heads + self.concat = concat + self.negative_slope = negative_slope + self.dropout = dropout + self.add_self_loops = add_self_loops + + self.lin = Linear(in_channels, heads * out_channels, bias=False) + + self.att = Parameter(torch.Tensor(1, heads, out_channels)) + + if bias and concat: + self.bias = Parameter(torch.Tensor(heads * out_channels)) + elif bias and not concat: + self.bias = Parameter(torch.Tensor(out_channels)) + else: + self.register_parameter('bias', None) + + self.reset_parameters() + + def reset_parameters(self): + glorot(self.lin.weight) + glorot(self.att) + zeros(self.bias) + + def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj): + H, C = self.heads, self.out_channels + + x = self.lin(x).view(-1, H, C) + alpha = (x * self.att).sum(dim=-1) + + if self.add_self_loops: + if isinstance(edge_index, Tensor): + num_nodes = x.size(0) + edge_index, _ = remove_self_loops(edge_index) + edge_index, _ = add_self_loops(edge_index, num_nodes=num_nodes) + elif isinstance(edge_index, SparseTensor): + edge_index = set_diag(edge_index) + + out = self.propagate(edge_index, x=x, + alpha=alpha) + + if self.concat: + out = out.view(-1, self.heads * self.out_channels) + else: + out = out.mean(dim=1) + + if self.bias is not None: + out += self.bias + + return out + + def message(self, x_j: Tensor, alpha_j: Tensor, alpha_i: OptTensor, + index: Tensor, ptr: OptTensor) -> Tensor: + alpha = alpha_j if alpha_i is None else alpha_j + alpha_i + alpha = F.leaky_relu(alpha, self.negative_slope) + alpha = softmax(alpha, index, ptr) + alpha = F.dropout(alpha, p=self.dropout, training=self.training) + return x_j * alpha.unsqueeze(-1) + + def __repr__(self): + return '{}({}, {}, heads={})'.format(self.__class__.__name__, + self.in_channels, + self.out_channels, self.heads) + +# class Zero(nn.Module): +# def __init__(self, indim, outdim) -> None: +# super().__init__() +# self.outdim = outdim +# self.zero = nn.Parameter(torch.tensor(0.), requires_grad=True) + +# def forward(self, x, edge_index): +# return torch.zeros(x.size(0), self.outdim).to(x.device) * self.zero + +class Zero(nn.Module): + def __init__(self, indim, outdim) -> None: + super().__init__() + self.outdim = outdim + self.ln = nn.Linear(1, 1) + + def forward(self, x, edge_index): + return 0. + +class Identity(nn.Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, x, edge_index): + return x + +class Linear(nn.Module): + def __init__(self, in_dim, out_dim): + super().__init__() + self.core = nn.Linear(in_dim, out_dim) + + def forward(self, x, *args): + return self.core(x) + + +agg_map={ + 'add':lambda dim,att_head=None,dropout=0,norm=False,skip_connect=False: AggAdd(dim,att_head,dropout,norm,skip_connect), + 'attn':lambda dim,att_head=None,dropout=0,norm=False,skip_connect=False: AggAttn(dim,att_head,dropout,norm,skip_connect), +} \ No newline at end of file diff --git a/autogl/module/nas/space/autoattend_space/ops1.py b/autogl/module/nas/space/autoattend_space/ops1.py new file mode 100644 index 0000000..45dffa5 --- /dev/null +++ b/autogl/module/nas/space/autoattend_space/ops1.py @@ -0,0 +1,19 @@ +from .operations import * + +OPS = { + 'ZERO': lambda indim, outdim, dropout, concat=False: Zero(indim, outdim), + 'IDEN': lambda indim, outdim, dropout, concat=False: Identity(), + 'GCN': lambda indim, outdim, dropout, concat=False: GCNConv(indim, outdim, add_self_loops=False), + 'SAGE-MEAN': lambda indim, outdim, dropout, concat=False: SAGEConv(indim, outdim), + 'GAT16': lambda indim, outdim, dropout, concat=False: GATConv(indim, outdim, dropout=dropout, heads=16, concat=False, add_self_loops=False) if not concat else GATConv(indim, outdim // 16, dropout=dropout, heads=16, concat=True, add_self_loops=False), + 'GAT2': lambda indim, outdim, dropout, concat=False: GATConv(indim, outdim, dropout=dropout, heads=2, concat=False, add_self_loops=False) if not concat else GATConv(indim, outdim // 2, dropout=dropout, heads=2, concat=True, add_self_loops=False), + 'GAT4': lambda indim, outdim, dropout, concat=False: GATConv(indim, outdim, dropout=dropout, heads=4, concat=False, add_self_loops=False) if not concat else GATConv(indim, outdim // 4, dropout=dropout, heads=4, concat=True, add_self_loops=False), + 'GAT8': lambda indim, outdim, dropout, concat=False: GATConv(indim, outdim, dropout=dropout, heads=8, concat=False, add_self_loops=False) if not concat else GATConv(indim, outdim // 8, dropout=dropout, heads=8, concat=True, add_self_loops=False), + 'GAT1': lambda indim, outdim, dropout, concat=False: GATConv(indim, outdim, dropout=dropout, heads=1, concat=False, add_self_loops=False), + 'LIN': lambda indim, outdim, dropout, concat=False: Linear(indim, outdim), + 'ARMA': lambda indim, outdim, dropout, concat=False: ARMAConv(indim, outdim), + 'CHEB': lambda indim, outdim, dropout, concat=False: ChebConv(indim, outdim, 2), + 'SGC': lambda indim, outdim, dropout, concat=False: SGConv(indim, outdim, add_self_loops=False) +} + +PRIMITIVES = list(OPS.keys()) \ No newline at end of file diff --git a/autogl/module/nas/space/autoattend_space/ops2.py b/autogl/module/nas/space/autoattend_space/ops2.py new file mode 100644 index 0000000..987cce3 --- /dev/null +++ b/autogl/module/nas/space/autoattend_space/ops2.py @@ -0,0 +1,13 @@ +from .operations import * +OPS = { + 'ZERO': lambda indim, outdim, head, dropout, concat=False: Zero(indim, outdim), + 'CONST': lambda indim, outdim, head, dropout, concat=False: GeoLayer(indim, outdim, head, concat, att_type='const', dropout=dropout), + 'GCN': lambda indim, outdim, head, dropout, concat=False: GeoLayer(indim, outdim, head, concat, att_type='gcn', dropout=dropout), + 'GAT': lambda indim, outdim, head, dropout, concat=False: GeoLayer(indim, outdim, head, concat, att_type='gat', dropout=dropout), + 'SYM': lambda indim, outdim, head, dropout, concat=False: GeoLayer(indim, outdim, head, concat, att_type='gat_sym', dropout=dropout), + 'COS': lambda indim, outdim, head, dropout, concat=False: GeoLayer(indim, outdim, head, concat, att_type='cos', dropout=dropout), + 'LIN': lambda indim, outdim, head, dropout, concat=False: GeoLayer(indim, outdim, head, concat, att_type='linear', dropout=dropout), + 'GENE': lambda indim, outdim, head, dropout, concat=False: GeoLayer(indim, outdim, head, concat, att_type='generalized_linear', dropout=dropout) +} + +PRIMITIVES = list(OPS.keys()) \ No newline at end of file diff --git a/autogl/module/nas/space/graph_nas_macro.py b/autogl/module/nas/space/graph_nas_macro.py index 3833977..68d4186 100644 --- a/autogl/module/nas/space/graph_nas_macro.py +++ b/autogl/module/nas/space/graph_nas_macro.py @@ -24,8 +24,10 @@ import inspect import sys from ..backend import * -import dgl -from dgl import function as fn +# import dgl +# from dgl import function as fn + + special_args = [ "edge_index", "edge_index_i", diff --git a/test/nas/node_classification.py b/test/nas/node_classification.py index 33c3acc..4875ab0 100644 --- a/test/nas/node_classification.py +++ b/test/nas/node_classification.py @@ -27,6 +27,7 @@ from autogl.module.nas.space.single_path import SinglePathNodeClassificationSpac from autogl.module.nas.space.graph_nas import GraphNasNodeClassificationSpace from autogl.module.nas.space.graph_nas_macro import GraphNasMacroNodeClassificationSpace from autogl.module.nas.estimator.one_shot import OneShotEstimator +from autogl.module.nas.space.autoattend import AutoAttendNodeClassificationSpace from autogl.module.nas.backend import bk_feat, bk_label from autogl.module.nas.algorithm import Darts, RL, GraphNasRL, Enas, RandomSearch import numpy as np @@ -81,66 +82,75 @@ if __name__ == "__main__": di = bk_feat(data).shape[1] do = len(np.unique(bk_label(data))) - print("Random search + singlepath ") - space = SinglePathNodeClassificationSpace().cuda() - space.instantiate(input_dim=di, output_dim=do) - esti = OneShotEstimator() - algo = RandomSearch(num_epochs=10) - model = algo.search(space, dataset, esti) - test_model(model, data, True) - - print("Random search + graphnas ") - space = GraphNasNodeClassificationSpace().cuda() - space.instantiate(input_dim=di, output_dim=do) - esti = OneShotEstimator() - algo = RandomSearch(num_epochs=10) - model = algo.search(space, dataset, esti) - test_model(model, data, True) + # print("Random search + singlepath ") + # space = SinglePathNodeClassificationSpace().cuda() + # space.instantiate(input_dim=di, output_dim=do) + # esti = OneShotEstimator() + # algo = RandomSearch(num_epochs=100) + # model = algo.search(space, dataset, esti) + # test_model(model, data, True) - print("rl + graphnas ") - space = GraphNasNodeClassificationSpace().cuda() - space.instantiate(input_dim=di, output_dim=do) - esti = OneShotEstimator() - algo = RL(num_epochs=10) - model = algo.search(space, dataset, esti) - test_model(model, data, True) - print("graphnasrl + graphnas ") - space = GraphNasNodeClassificationSpace().cuda() + print("Random search + AutoAttend ") + space = AutoAttendNodeClassificationSpace().cuda() space.instantiate(input_dim=di, output_dim=do) esti = OneShotEstimator() - algo = GraphNasRL(num_epochs=10) + algo = RandomSearch(num_epochs=100) model = algo.search(space, dataset, esti) test_model(model, data, True) - print("darts + graphnas ") - space = GraphNasNodeClassificationSpace(con_ops=['concat']).cuda() - space.instantiate(input_dim=di, output_dim=do) - esti = OneShotEstimator() - algo = Darts(num_epochs=10) - model = algo.search(space, dataset, esti) - test_model(model, data, True) + # print("Random search + graphnas ") + # space = GraphNasNodeClassificationSpace().cuda() + # space.instantiate(input_dim=di, output_dim=do) + # esti = OneShotEstimator() + # algo = RandomSearch(num_epochs=10) + # model = algo.search(space, dataset, esti) + # test_model(model, data, True) + + # print("rl + graphnas ") + # space = GraphNasNodeClassificationSpace().cuda() + # space.instantiate(input_dim=di, output_dim=do) + # esti = OneShotEstimator() + # algo = RL(num_epochs=10) + # model = algo.search(space, dataset, esti) + # test_model(model, data, True) + + # print("graphnasrl + graphnas ") + # space = GraphNasNodeClassificationSpace().cuda() + # space.instantiate(input_dim=di, output_dim=do) + # esti = OneShotEstimator() + # algo = GraphNasRL(num_epochs=10) + # model = algo.search(space, dataset, esti) + # test_model(model, data, True) + + # print("darts + graphnas ") + # space = GraphNasNodeClassificationSpace(con_ops=['concat']).cuda() + # space.instantiate(input_dim=di, output_dim=do) + # esti = OneShotEstimator() + # algo = Darts(num_epochs=10) + # model = algo.search(space, dataset, esti) + # test_model(model, data, True) - print("darts + singlepath ") - space = SinglePathNodeClassificationSpace().cuda() - space.instantiate(input_dim=di, output_dim=do) - esti = OneShotEstimator() - algo = Darts(num_epochs=10) - model = algo.search(space, dataset, esti) - test_model(model, data, True) - - print("Random search + graphnas macro") - space = GraphNasMacroNodeClassificationSpace().cuda() - space.instantiate(input_dim=di, output_dim=do) - esti = OneShotEstimator() - algo = RandomSearch(num_epochs=10) - model = algo.search(space, dataset, esti) - test_model(model, data, True) - - print("RL + graphnas macro") - space = GraphNasMacroNodeClassificationSpace().cuda() - space.instantiate(input_dim=di, output_dim=do) - esti = OneShotEstimator() - algo = RL(num_epochs=10) - model = algo.search(space, dataset, esti) - test_model(model, data, True) + # print("darts + singlepath ") + # space = SinglePathNodeClassificationSpace().cuda() + # space.instantiate(input_dim=di, output_dim=do) + # esti = OneShotEstimator() + # algo = Darts(num_epochs=10) + # model = algo.search(space, dataset, esti) + # test_model(model, data, True) + + # print("Random search + graphnas macro") + # space = GraphNasMacroNodeClassificationSpace().cuda() + # space.instantiate(input_dim=di, output_dim=do) + # esti = OneShotEstimator() + # algo = RandomSearch(num_epochs=10) + # model = algo.search(space, dataset, esti) + # test_model(model, data, True) + + # print("RL + graphnas macro") + # space = GraphNasMacroNodeClassificationSpace().cuda() + # space.instantiate(input_dim=di, output_dim=do) + # esti = OneShotEstimator() + # algo = RL(num_epochs=10) + # model = algo.search(space, dataset, esti) + # test_model(model, data, True) From 5f49726662236db8e2fdbecae19e9299ba7637df Mon Sep 17 00:00:00 2001 From: wondergo2017 Date: Fri, 17 Dec 2021 23:36:54 +0800 Subject: [PATCH 06/14] fix autoattend aggattn propogate;zeroconv;Multilayerchoice mutable --- autogl/module/nas/algorithm/random_search.py | 1 + autogl/module/nas/space/autoattend.py | 49 ++++++++++--------- .../nas/space/autoattend_space/operations.py | 48 ++++++++++-------- autogl/module/nas/space/base.py | 19 ++++--- autogl/module/nas/utils.py | 2 + test/nas/node_classification.py | 2 +- 6 files changed, 71 insertions(+), 50 deletions(-) diff --git a/autogl/module/nas/algorithm/random_search.py b/autogl/module/nas/algorithm/random_search.py index f704dbd..4670406 100644 --- a/autogl/module/nas/algorithm/random_search.py +++ b/autogl/module/nas/algorithm/random_search.py @@ -73,6 +73,7 @@ class RandomSearch(BaseNAS): 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) + print("randomsearch",arch) return arch def sample(self): diff --git a/autogl/module/nas/space/autoattend.py b/autogl/module/nas/space/autoattend.py index a33b375..2ec99d4 100644 --- a/autogl/module/nas/space/autoattend.py +++ b/autogl/module/nas/space/autoattend.py @@ -22,9 +22,11 @@ from .autoattend_space.operations import agg_map OPS = [OPS1, OPS2] from nni.nas.pytorch.mutables import Mutable -class MultiLayerChoice(Mutable): +# class MultiLayerChoice(Mutable,nn.Module): +class MultiLayerChoice(nn.Module): def __init__(self, choices, layer, key): - super(MultiLayerChoice, self).__init__(key) + # Mutable.__init__(self,key) + nn.Module.__init__(self) self.order = layer self.choices = choices @@ -101,14 +103,12 @@ class AutoAttendNodeClassificationSpace(BaseSpace): self.preproc0 = nn.Linear(self.input_dim, self.hidden_dim) # self.preproc1 = nn.Linear(self.input_dim, self.hidden_dim) node_labels = [] - - # stem path for layer in range(1, self.layer_number+1): + # stem path key = f"stem_{layer}" self._set_layer_choice(layer, key) - # side path - for layer in range(1, self.layer_number+1): + # side path choices = [] key = f"side_{layer}" for i in range(2): @@ -120,19 +120,18 @@ class AutoAttendNodeClassificationSpace(BaseSpace): # input key = f"in_{layer}" - self._set_input_choice(key, - layer, choose_from=node_labels, n_chosen=1, return_mask=False) - - # agg - for layer in range(1, self.layer_number + 1): - key = f"agg_{layer}" - self._set_agg_choice(layer, key=key) + # self._set_input_choice(key,layer, choose_from=node_labels, n_chosen=1, return_mask=False) + self._set_input_choice(key,layer, n_candidates=len(node_labels), n_chosen=1, return_mask=False) + # agg + key = f"agg_{layer}" + self._set_agg_choice(layer, key=key) + self._initialized = True self.classifier2 = nn.Linear(self.hidden_dim, self.output_dim) - print(self) + # print(self) def _set_agg_choice(self,layer,key): ops=[self.agg_map(op, self.hidden_dim,self.head,self.dropout)for op in self.agg_ops] choice = self.setLayerChoice( @@ -171,17 +170,20 @@ class AutoAttendNodeClassificationSpace(BaseSpace): side_outs = [] stem_outs = [] - input = prev_ + input = prev_ # torch.Size([2708, 64]) for layer in range(1, self.layer_number + 1): + print(f'start {layer}') # do layer choice for stem op = getattr(self, f"stem_{layer}") - print(op) + print(f"stem op {op}") stem_out = bk_gconv(op, data, input) - stem_out = self.act(stem_out) + stem_out = self.act(stem_out) #torch.Size([2708, 64]) # do double layer choice for sides - op = getattr(self, f'side_{layer}') + op = getattr(self, f'side_{layer}') + print(f"side op {op}") + print("input",input.shape) side_out = bk_gconv(op, data, input) - side_out = self.act(side_out) + side_out = self.act(side_out) # torch.Size([2, 2708, 64]) stem_outs.append(stem_out) side_outs.append(side_out) @@ -191,14 +193,17 @@ class AutoAttendNodeClassificationSpace(BaseSpace): input = [stem_outs[-1], side_selected] # do agg in [add , attn] agg = getattr(self, f"agg_{layer}") - print(layer,input) + + print(stem_out.shape,side_out.shape,agg) + input = bk_gconv(agg,data,input) + print(f'done {layer}') x = self.classifier2(input) return F.log_softmax(x, dim=1) def parse_model(self, selection, device) -> BaseModel: - for i in range(1, self.layer_number + 1): - selection[f'side_{i}'] = None + # for i in range(1, self.layer_number + 1): + # selection[f'side_{i}'] = None # return AutoGCN(self.input_dim, self.output_dim, device) return self.wrap(device).fix(selection) diff --git a/autogl/module/nas/space/autoattend_space/operations.py b/autogl/module/nas/space/autoattend_space/operations.py index 74b7f7f..18da1f3 100644 --- a/autogl/module/nas/space/autoattend_space/operations.py +++ b/autogl/module/nas/space/autoattend_space/operations.py @@ -63,11 +63,12 @@ class AggAttn(MessagePassing) : # use dot-product attn x1,x2,x3=x[0],x[1][0],x[1][1] # q,k,v skip_connect,norm=self.skip_connect,self.norm - + print('AggAttn') + print(x1.shape,x2.shape,x3.shape) if not skip_connect and not norm: - return self.propagate(edge_index, x1=x1, x2=x2, x3=x3).view(x1.size(0), -1) + return self.propagate(edge_index, x1=x1, x2=x2, x3=x3) - x = self.propagate(edge_index, x1=x1, x2=x2, x3=x3).view(x1.size(0), -1) + x = self.propagate(edge_index, x1=x1, x2=x2, x3=x3) if not norm: return x @@ -76,15 +77,20 @@ class AggAttn(MessagePassing) : return self.ln_attn(x + x1) def message(self, x2_j, x1_i, x3_j, index, ptr): - # x1: query, x2: key, x3: value + print('message') + print(index.shape) + # x1: query, x2: key, x3: value # torch.Size([10556, 64]) ,index torch.Size([10556]) node, dim = x1_i.size() dim_att = dim // self.att_head - x2_j = x2_j.view(node, self.att_head, dim_att) - x1_i = x1_i.view(node, self.att_head, dim_att) - attn = (x2_j * x1_i).sum(dim=-1) / np.sqrt(dim_att) - attn = softmax(attn, index, ptr) - attn = F.dropout(attn, p=self.dropout, training=self.training) - return x3_j.view(node, self.att_head, dim_att) * attn.unsqueeze(-1) + x2_j = x2_j.view(node, self.att_head, dim_att) #torch.Size([10556, 8, 8]) + x1_i = x1_i.view(node, self.att_head, dim_att) #torch.Size([10556, 8, 8]) + attn = (x2_j * x1_i).sum(dim=-1) / np.sqrt(dim_att) #torch.Size([10556, 8]) + attn = softmax(attn, index, ptr) #torch.Size([10556, 8]) + attn = F.dropout(attn, p=self.dropout, training=self.training) #torch.Size([10556, 8]) + out=x3_j.view(node, self.att_head, dim_att) * attn.unsqueeze(-1) + out=out.view(-1,dim) + print(out.shape) + return out class GATConv2(MessagePassing): _alpha: OptTensor @@ -161,23 +167,23 @@ class GATConv2(MessagePassing): self.in_channels, self.out_channels, self.heads) -# class Zero(nn.Module): -# def __init__(self, indim, outdim) -> None: -# super().__init__() -# self.outdim = outdim -# self.zero = nn.Parameter(torch.tensor(0.), requires_grad=True) - -# def forward(self, x, edge_index): -# return torch.zeros(x.size(0), self.outdim).to(x.device) * self.zero - class Zero(nn.Module): def __init__(self, indim, outdim) -> None: super().__init__() self.outdim = outdim - self.ln = nn.Linear(1, 1) + self.zero = nn.Parameter(torch.tensor(0.), requires_grad=True) def forward(self, x, edge_index): - return 0. + return torch.zeros(x.size(0), self.outdim).to(x.device) * self.zero + +# class Zero(nn.Module): +# def __init__(self, indim, outdim) -> None: +# super().__init__() +# self.outdim = outdim +# self.ln = nn.Linear(1, 1) + +# def forward(self, x, edge_index): +# return 0. class Identity(nn.Module): def __init__(self) -> None: diff --git a/autogl/module/nas/space/base.py b/autogl/module/nas/space/base.py index 8ff7901..ec4d691 100644 --- a/autogl/module/nas/space/base.py +++ b/autogl/module/nas/space/base.py @@ -10,7 +10,6 @@ from ....utils import get_logger from ..utils import get_hardware_aware_metric - class OrderedMutable: """ An abstract class with order, enabling to sort mutables with a certain rank. @@ -30,7 +29,8 @@ class OrderedLayerChoice(OrderedMutable, mutables.LayerChoice): 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) + mutables.LayerChoice.__init__( + self, op_candidates, reduction, return_mask, key) class OrderedInputChoice(OrderedMutable, mutables.InputChoice): @@ -106,7 +106,8 @@ class BoxModel(BaseModel): self._model = space_model.to(device) self.num_features = self._model.input_dim self.num_classes = self._model.output_dim - self.params = {"num_class": self.num_classes, "features_num": self.num_features} + self.params = {"num_class": self.num_classes, + "features_num": self.num_features} self.device = device self.selection = None @@ -140,12 +141,17 @@ class BoxModel(BaseModel): ret_self = deepcopy(self) ret_self._model.instantiate() if ret_self.selection: - apply_fixed_architecture(ret_self._model, ret_self.selection, verbose=False) + apply_fixed_architecture( + ret_self._model, ret_self.selection, verbose=False) ret_self.to(self.device) return ret_self def __repr__(self) -> str: - return str({'parameter': get_hardware_aware_metric(self.model, 'parameter')}) + return str( + {'parameter': get_hardware_aware_metric(self.model, 'parameter'), + 'model': self.model, + 'selection': self.selection + }) @property def model(self): @@ -221,7 +227,8 @@ class BaseSpace(nn.Module): key = f"default_key_{self._default_key}" self._default_key += 1 orikey = key - layer = OrderedLayerChoice(order, op_candidates, reduction, return_mask, orikey) + layer = OrderedLayerChoice( + order, op_candidates, reduction, return_mask, orikey) return layer def setInputChoice( diff --git a/autogl/module/nas/utils.py b/autogl/module/nas/utils.py index 45f546f..31d3793 100644 --- a/autogl/module/nas/utils.py +++ b/autogl/module/nas/utils.py @@ -62,6 +62,8 @@ class PathSamplingLayerChoice(nn.Module): def mask(self): return _get_mask(self.sampled, len(self)) + def __repr__(self): + return f"PathSamplingLayerChoice(op_names={self.op_names}, chosen={self.sampled})" class PathSamplingInputChoice(nn.Module): """ diff --git a/test/nas/node_classification.py b/test/nas/node_classification.py index 4875ab0..940c020 100644 --- a/test/nas/node_classification.py +++ b/test/nas/node_classification.py @@ -95,7 +95,7 @@ if __name__ == "__main__": space = AutoAttendNodeClassificationSpace().cuda() space.instantiate(input_dim=di, output_dim=do) esti = OneShotEstimator() - algo = RandomSearch(num_epochs=100) + algo = RandomSearch(num_epochs=10) model = algo.search(space, dataset, esti) test_model(model, data, True) From c382810195864c67c854513dcea9589acc193d33 Mon Sep 17 00:00:00 2001 From: wondergo2017 Date: Sat, 18 Dec 2021 00:30:09 +0800 Subject: [PATCH 07/14] fix deepcopy mutable problem;format codes --- autogl/module/nas/algorithm/random_search.py | 1 - autogl/module/nas/space/autoattend.py | 90 +++++++------------ .../nas/space/autoattend_space/operations.py | 85 ++++++++++-------- .../module/nas/space/autoattend_space/ops1.py | 2 +- .../module/nas/space/autoattend_space/ops2.py | 2 +- autogl/module/nas/space/base.py | 1 + test/nas/node_classification.py | 1 + 7 files changed, 85 insertions(+), 97 deletions(-) diff --git a/autogl/module/nas/algorithm/random_search.py b/autogl/module/nas/algorithm/random_search.py index 4670406..f704dbd 100644 --- a/autogl/module/nas/algorithm/random_search.py +++ b/autogl/module/nas/algorithm/random_search.py @@ -73,7 +73,6 @@ class RandomSearch(BaseNAS): 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) - print("randomsearch",arch) return arch def sample(self): diff --git a/autogl/module/nas/space/autoattend.py b/autogl/module/nas/space/autoattend.py index 2ec99d4..650726f 100644 --- a/autogl/module/nas/space/autoattend.py +++ b/autogl/module/nas/space/autoattend.py @@ -21,30 +21,13 @@ from .autoattend_space.ops2 import OPS as OPS2 from .autoattend_space.operations import agg_map OPS = [OPS1, OPS2] -from nni.nas.pytorch.mutables import Mutable -# class MultiLayerChoice(Mutable,nn.Module): -class MultiLayerChoice(nn.Module): - def __init__(self, choices, layer, key): - # Mutable.__init__(self,key) - nn.Module.__init__(self) - self.order = layer - self.choices = choices - - def forward(self, *args, **kwargs): - outs = [] - for i in range(len(self.choices)): - out = self.choices[i](*args, **kwargs) - outs.append(out) - outs = torch.stack(outs, dim=0) - return outs - @register_nas_space("autoattend") class AutoAttendNodeClassificationSpace(BaseSpace): def __init__( self, hidden_dim: _typ.Optional[int] = 64, - layer_number: _typ.Optional[int] = 2, + layer_number: _typ.Optional[int] = 5, dropout: _typ.Optional[float] = 0.9, input_dim: _typ.Optional[int] = None, output_dim: _typ.Optional[int] = None, @@ -53,7 +36,7 @@ class AutoAttendNodeClassificationSpace(BaseSpace): ] = None, act_op="tanh", head=8, - agg_ops=['add','attn'] + agg_ops=['add', 'attn'] ): super().__init__() self.layer_number = layer_number @@ -92,16 +75,15 @@ class AutoAttendNodeClassificationSpace(BaseSpace): self.act = act_map(self.act_op) self.head = head or self.head self.ops_type = ops_type or self.ops_type - self.agg_ops=agg_ops or self.agg_ops + self.agg_ops = agg_ops or self.agg_ops PRIMITIVES = list(OPS[self.ops_type].keys()) self.gnn_map = lambda x, * \ args, **kwargs: OPS[self.ops_type][x](*args, **kwargs) self.gnn_ops = self.gnn_ops or PRIMITIVES self.agg_map = lambda x, * \ - args, **kwargs: agg_map[x](*args, **kwargs) - # self.con_ops = con_ops or self.con_ops + args, **kwargs: agg_map[x](*args, **kwargs) self.preproc0 = nn.Linear(self.input_dim, self.hidden_dim) - # self.preproc1 = nn.Linear(self.input_dim, self.hidden_dim) + node_labels = [] for layer in range(1, self.layer_number+1): # stem path @@ -109,31 +91,29 @@ class AutoAttendNodeClassificationSpace(BaseSpace): self._set_layer_choice(layer, key) # side path - choices = [] key = f"side_{layer}" for i in range(2): sub_key = f"{key}_{i}" - choice = self._set_layer_choice(layer, sub_key) - choices.append(choice) - setattr(self, key, MultiLayerChoice(choices, layer, key)) + self._set_layer_choice(layer, sub_key) node_labels.append(key) # input key = f"in_{layer}" - # self._set_input_choice(key,layer, choose_from=node_labels, n_chosen=1, return_mask=False) - self._set_input_choice(key,layer, n_candidates=len(node_labels), n_chosen=1, return_mask=False) + # self._set_input_choice(key,layer, choose_from=node_labels, n_chosen=1, return_mask=False) + self._set_input_choice(key, layer, n_candidates=len( + node_labels), n_chosen=1, return_mask=False) # agg key = f"agg_{layer}" - self._set_agg_choice(layer, key=key) - + self._set_agg_choice(layer, key=key) + self._initialized = True self.classifier2 = nn.Linear(self.hidden_dim, self.output_dim) - # print(self) - def _set_agg_choice(self,layer,key): - ops=[self.agg_map(op, self.hidden_dim,self.head,self.dropout)for op in self.agg_ops] + def _set_agg_choice(self, layer, key): + ops = [self.agg_map(op, self.hidden_dim, self.head, + self.dropout)for op in self.agg_ops] choice = self.setLayerChoice( layer, ops, @@ -141,11 +121,14 @@ class AutoAttendNodeClassificationSpace(BaseSpace): ) setattr(self, key, choice) return choice + def _set_layer_choice(self, layer, key): - if self.ops_type==0: - ops=[self.gnn_map(op, self.hidden_dim, self.hidden_dim, self.dropout)for op in self.gnn_ops] - elif self.ops_type==1: - ops=[self.gnn_map(op, self.hidden_dim, self.hidden_dim, self.head, self.dropout)for op in self.gnn_ops] + if self.ops_type == 0: + ops = [self.gnn_map( + op, self.hidden_dim, self.hidden_dim, self.dropout)for op in self.gnn_ops] + elif self.ops_type == 1: + ops = [self.gnn_map(op, self.hidden_dim, self.hidden_dim, + self.head, self.dropout)for op in self.gnn_ops] choice = self.setLayerChoice( layer, ops, @@ -154,7 +137,7 @@ class AutoAttendNodeClassificationSpace(BaseSpace): setattr(self, key, choice) return choice - def _set_input_choice(self, key, layer,**kwargs): + def _set_input_choice(self, key, layer, **kwargs): setattr(self, key, self.setInputChoice( @@ -170,20 +153,21 @@ class AutoAttendNodeClassificationSpace(BaseSpace): side_outs = [] stem_outs = [] - input = prev_ # torch.Size([2708, 64]) + input = prev_ for layer in range(1, self.layer_number + 1): - print(f'start {layer}') # do layer choice for stem op = getattr(self, f"stem_{layer}") - print(f"stem op {op}") stem_out = bk_gconv(op, data, input) - stem_out = self.act(stem_out) #torch.Size([2708, 64]) + stem_out = self.act(stem_out) + # do double layer choice for sides - op = getattr(self, f'side_{layer}') - print(f"side op {op}") - print("input",input.shape) - side_out = bk_gconv(op, data, input) - side_out = self.act(side_out) # torch.Size([2, 2708, 64]) + side_out_list = [] + for i in range(2): + op = getattr(self, f'side_{layer}_{i}') + side_out = bk_gconv(op, data, input) + side_out = self.act(side_out) # torch.Size([2, 2708, 64]) + side_out_list.append(side_out) + side_out = torch.stack(side_out_list, dim=0) stem_outs.append(stem_out) side_outs.append(side_out) @@ -191,19 +175,13 @@ class AutoAttendNodeClassificationSpace(BaseSpace): # select input [x1,x2,x3] from side1,side2,stem side_selected = getattr(self, f"in_{layer}")(side_outs) input = [stem_outs[-1], side_selected] + # do agg in [add , attn] agg = getattr(self, f"agg_{layer}") - - print(stem_out.shape,side_out.shape,agg) - - input = bk_gconv(agg,data,input) - print(f'done {layer}') + input = bk_gconv(agg, data, input) x = self.classifier2(input) return F.log_softmax(x, dim=1) def parse_model(self, selection, device) -> BaseModel: - # for i in range(1, self.layer_number + 1): - # selection[f'side_{i}'] = None - # return AutoGCN(self.input_dim, self.output_dim, device) return self.wrap(device).fix(selection) diff --git a/autogl/module/nas/space/autoattend_space/operations.py b/autogl/module/nas/space/autoattend_space/operations.py index 18da1f3..73b1b4f 100644 --- a/autogl/module/nas/space/autoattend_space/operations.py +++ b/autogl/module/nas/space/autoattend_space/operations.py @@ -1,4 +1,4 @@ - + from torch_geometric.nn import MessagePassing from torch_geometric.utils import softmax import torch @@ -31,43 +31,47 @@ from torch_scatter import scatter_add import numpy as np from ..graph_nas_macro import GeoLayer + class AggAdd(nn.Module): - def __init__(self, dim, att_head, dropout=0,norm=False,skip_connect=False,*args,**kwargs): + def __init__(self, dim, att_head, dropout=0, norm=False, skip_connect=False, *args, **kwargs): super(AggAdd, self).__init__() self.dropout = dropout - self.ln_add = nn.BatchNorm1d(dim, track_running_stats=True, affine=True) - self.norm=norm - self.skip_connect=skip_connect + self.ln_add = nn.BatchNorm1d( + dim, track_running_stats=True, affine=True) + self.norm = norm + self.skip_connect = skip_connect - def forward(self,x,edge_index,*args,**kwargs): + def forward(self, x, edge_index, *args, **kwargs): # x=[x_stem,[x_sides]] - norm=self.norm - x1,x2,x3=x[0],x[1][0],x[1][1] + norm = self.norm + x1, x2, x3 = x[0], x[1][0], x[1][1] if norm: return self.ln_add(x1 + x2) else: return x1 + x2 -class AggAttn(MessagePassing) : - def __init__(self, dim, att_head,dropout=0,norm=False,skip_connect=False,*args,**kwargs): + +class AggAttn(MessagePassing): + def __init__(self, dim, att_head, dropout=0, norm=False, skip_connect=False, *args, **kwargs): super(AggAttn, self).__init__() self.dropout = dropout - self.att_head=att_head - self.ln_attn = nn.BatchNorm1d(dim, track_running_stats=True, affine=True) - self.norm=norm - self.skip_connect=skip_connect + self.att_head = att_head + self.ln_attn = nn.BatchNorm1d( + dim, track_running_stats=True, affine=True) + self.norm = norm + self.skip_connect = skip_connect + def __repr__(self) -> str: return 'AggAttn(att_head={}, dropout={})'.format(self.att_head, self.dropout) - def forward(self,x,edge_index,*args,**kwargs): + + def forward(self, x, edge_index, *args, **kwargs): # x=[x_stem,[x_sides]] # use dot-product attn - x1,x2,x3=x[0],x[1][0],x[1][1] # q,k,v - skip_connect,norm=self.skip_connect,self.norm - print('AggAttn') - print(x1.shape,x2.shape,x3.shape) + x1, x2, x3 = x[0], x[1][0], x[1][1] # q,k,v + skip_connect, norm = self.skip_connect, self.norm if not skip_connect and not norm: return self.propagate(edge_index, x1=x1, x2=x2, x3=x3) - + x = self.propagate(edge_index, x1=x1, x2=x2, x3=x3) if not norm: return x @@ -77,21 +81,23 @@ class AggAttn(MessagePassing) : return self.ln_attn(x + x1) def message(self, x2_j, x1_i, x3_j, index, ptr): - print('message') - print(index.shape) # x1: query, x2: key, x3: value # torch.Size([10556, 64]) ,index torch.Size([10556]) node, dim = x1_i.size() dim_att = dim // self.att_head - x2_j = x2_j.view(node, self.att_head, dim_att) #torch.Size([10556, 8, 8]) - x1_i = x1_i.view(node, self.att_head, dim_att) #torch.Size([10556, 8, 8]) - attn = (x2_j * x1_i).sum(dim=-1) / np.sqrt(dim_att) #torch.Size([10556, 8]) - attn = softmax(attn, index, ptr) #torch.Size([10556, 8]) - attn = F.dropout(attn, p=self.dropout, training=self.training) #torch.Size([10556, 8]) - out=x3_j.view(node, self.att_head, dim_att) * attn.unsqueeze(-1) - out=out.view(-1,dim) - print(out.shape) + # torch.Size([10556, 8, 8]) + x2_j = x2_j.view(node, self.att_head, dim_att) + # torch.Size([10556, 8, 8]) + x1_i = x1_i.view(node, self.att_head, dim_att) + attn = (x2_j * x1_i).sum(dim=-1) / \ + np.sqrt(dim_att) # torch.Size([10556, 8]) + attn = softmax(attn, index, ptr) # torch.Size([10556, 8]) + # torch.Size([10556, 8]) + attn = F.dropout(attn, p=self.dropout, training=self.training) + out = x3_j.view(node, self.att_head, dim_att) * attn.unsqueeze(-1) + out = out.view(-1, dim) return out + class GATConv2(MessagePassing): _alpha: OptTensor @@ -167,12 +173,13 @@ class GATConv2(MessagePassing): self.in_channels, self.out_channels, self.heads) + class Zero(nn.Module): def __init__(self, indim, outdim) -> None: super().__init__() self.outdim = outdim self.zero = nn.Parameter(torch.tensor(0.), requires_grad=True) - + def forward(self, x, edge_index): return torch.zeros(x.size(0), self.outdim).to(x.device) * self.zero @@ -181,27 +188,29 @@ class Zero(nn.Module): # super().__init__() # self.outdim = outdim # self.ln = nn.Linear(1, 1) - + # def forward(self, x, edge_index): # return 0. + class Identity(nn.Module): def __init__(self) -> None: super().__init__() - + def forward(self, x, edge_index): return x + class Linear(nn.Module): def __init__(self, in_dim, out_dim): super().__init__() self.core = nn.Linear(in_dim, out_dim) - + def forward(self, x, *args): return self.core(x) -agg_map={ - 'add':lambda dim,att_head=None,dropout=0,norm=False,skip_connect=False: AggAdd(dim,att_head,dropout,norm,skip_connect), - 'attn':lambda dim,att_head=None,dropout=0,norm=False,skip_connect=False: AggAttn(dim,att_head,dropout,norm,skip_connect), -} \ No newline at end of file +agg_map = { + 'add': lambda dim, att_head=None, dropout=0, norm=False, skip_connect=False: AggAdd(dim, att_head, dropout, norm, skip_connect), + 'attn': lambda dim, att_head=None, dropout=0, norm=False, skip_connect=False: AggAttn(dim, att_head, dropout, norm, skip_connect), +} diff --git a/autogl/module/nas/space/autoattend_space/ops1.py b/autogl/module/nas/space/autoattend_space/ops1.py index 45dffa5..9fc8c21 100644 --- a/autogl/module/nas/space/autoattend_space/ops1.py +++ b/autogl/module/nas/space/autoattend_space/ops1.py @@ -16,4 +16,4 @@ OPS = { 'SGC': lambda indim, outdim, dropout, concat=False: SGConv(indim, outdim, add_self_loops=False) } -PRIMITIVES = list(OPS.keys()) \ No newline at end of file +PRIMITIVES = list(OPS.keys()) diff --git a/autogl/module/nas/space/autoattend_space/ops2.py b/autogl/module/nas/space/autoattend_space/ops2.py index 987cce3..49b2e18 100644 --- a/autogl/module/nas/space/autoattend_space/ops2.py +++ b/autogl/module/nas/space/autoattend_space/ops2.py @@ -10,4 +10,4 @@ OPS = { 'GENE': lambda indim, outdim, head, dropout, concat=False: GeoLayer(indim, outdim, head, concat, att_type='generalized_linear', dropout=dropout) } -PRIMITIVES = list(OPS.keys()) \ No newline at end of file +PRIMITIVES = list(OPS.keys()) diff --git a/autogl/module/nas/space/base.py b/autogl/module/nas/space/base.py index ec4d691..7807174 100644 --- a/autogl/module/nas/space/base.py +++ b/autogl/module/nas/space/base.py @@ -318,6 +318,7 @@ class CleanFixedArchitecture(FixedArchitecture): prefix : str Module name under global namespace. """ + if module is None: module = self.model for name, mutable in module.named_children(): diff --git a/test/nas/node_classification.py b/test/nas/node_classification.py index 940c020..a60bb44 100644 --- a/test/nas/node_classification.py +++ b/test/nas/node_classification.py @@ -97,6 +97,7 @@ if __name__ == "__main__": esti = OneShotEstimator() algo = RandomSearch(num_epochs=10) model = algo.search(space, dataset, esti) + print(model) test_model(model, data, True) # print("Random search + graphnas ") From 7728fffa1fb62fbec717874ce4f2b4c3a3c18999 Mon Sep 17 00:00:00 2001 From: generall Date: Fri, 24 Dec 2021 11:20:16 +0800 Subject: [PATCH 08/14] update nas tutorial --- docs/docfile/tutorial/t_nas.rst | 20 +++++++++++++++++++- resources/nas.svg | 1 + 2 files changed, 20 insertions(+), 1 deletion(-) create mode 100644 resources/nas.svg diff --git a/docs/docfile/tutorial/t_nas.rst b/docs/docfile/tutorial/t_nas.rst index 10ed15d..6fbaad8 100644 --- a/docs/docfile/tutorial/t_nas.rst +++ b/docs/docfile/tutorial/t_nas.rst @@ -4,7 +4,25 @@ Neural Architecture Search ============================ We support different neural architecture search algorithm in variant search space. -To be more flexible, we modulize NAS process with three part: algorithm, space and estimator. +Neural architecture search is usually constructed by three modules: search space, search strategy and estimation strategy. + +The search space describes all possible architectures to be searched. There are mainly two parts of the space formulated, the operations(e.g. GCNconv, GATconv) and the input-ouput realations. +A large space may have better optimal architecture but demands more effect to explore. +Human knowledge can help to design a reasonable search space to reduce the efforts of search strategy. + +The search strategy controls how to explore the search sapce. +It encompasses the classical exploration-exploitation trade-off since. +On the one hand, it is desirable to find well-performing architectures quickly, +while on the other hand, premature convergence to a region of suboptimal architectures should be avoided. + +The estimation strategy gives the performance of certain architectures when it is explored. +The simplest option is to perform a standard training and validation of the architecture on data. +Since there are lots of architectures need estimating in the whole searching process, estimation strategy is desired to be very efficient to save computational resources. + +.. image:: ../resources/nas.svg + :align: center + +To be more flexible, we modulize NAS process with three part: algorithm, space and estimator, corresponding to the three module search space, search strategy and estimation strategy. Different models in different parts can be composed in some certain constrains. If you want to design your own NAS process, you can change any of those parts according to your demand. diff --git a/resources/nas.svg b/resources/nas.svg new file mode 100644 index 0000000..a2168ca --- /dev/null +++ b/resources/nas.svg @@ -0,0 +1 @@ +NAS frameworkSearch SpaceSearch StrategyEstimation StrategyspacealgorithmestimatorAutoGLclass \ No newline at end of file From e495b6f7aef3e897681f6cb80ce2b1f4e5a7a57c Mon Sep 17 00:00:00 2001 From: generall Date: Sat, 25 Dec 2021 11:42:23 +0800 Subject: [PATCH 09/14] fix default cuda --- autogl/module/nas/algorithm/darts.py | 2 +- autogl/module/nas/algorithm/enas.py | 2 +- autogl/module/nas/algorithm/gasso.py | 2 +- autogl/module/nas/algorithm/random_search.py | 2 +- autogl/module/nas/algorithm/rl.py | 4 ++-- autogl/module/nas/space/base.py | 4 ++-- autogl/module/nas/space/gasso.py | 7 ++++++- 7 files changed, 14 insertions(+), 9 deletions(-) diff --git a/autogl/module/nas/algorithm/darts.py b/autogl/module/nas/algorithm/darts.py index 87627e7..4d85292 100644 --- a/autogl/module/nas/algorithm/darts.py +++ b/autogl/module/nas/algorithm/darts.py @@ -103,7 +103,7 @@ class Darts(BaseNAS): model_wd=5e-4, arch_lr=3e-4, arch_wd=1e-3, - device="cuda", + device="auto", ): super().__init__(device=device) self.num_epochs = num_epochs diff --git a/autogl/module/nas/algorithm/enas.py b/autogl/module/nas/algorithm/enas.py index 36fac6e..d2c8dc3 100644 --- a/autogl/module/nas/algorithm/enas.py +++ b/autogl/module/nas/algorithm/enas.py @@ -79,7 +79,7 @@ class Enas(BaseNAS): model_lr=5e-3, model_wd=5e-4, disable_progress=True, - device="cuda", + device="auto", ): super().__init__(device) self.device = device diff --git a/autogl/module/nas/algorithm/gasso.py b/autogl/module/nas/algorithm/gasso.py index 42345cc..fc50ed6 100644 --- a/autogl/module/nas/algorithm/gasso.py +++ b/autogl/module/nas/algorithm/gasso.py @@ -34,7 +34,7 @@ class Gasso(BaseNAS): arch_lr = 0.03, stru_lr = 0.04, lamb = 0.6, - device="cuda", + device="auto", ): super().__init__(device=device) self.device = device diff --git a/autogl/module/nas/algorithm/random_search.py b/autogl/module/nas/algorithm/random_search.py index f704dbd..8abd292 100644 --- a/autogl/module/nas/algorithm/random_search.py +++ b/autogl/module/nas/algorithm/random_search.py @@ -35,7 +35,7 @@ class RandomSearch(BaseNAS): Control whether show the progress bar. """ - def __init__(self, device="cuda", num_epochs=400, disable_progress=False, hardware_metric_limit=None): + def __init__(self, device="auto", num_epochs=400, disable_progress=False, hardware_metric_limit=None): super().__init__(device) self.num_epochs = num_epochs self.disable_progress = disable_progress diff --git a/autogl/module/nas/algorithm/rl.py b/autogl/module/nas/algorithm/rl.py index 2f8e9a8..d927c26 100644 --- a/autogl/module/nas/algorithm/rl.py +++ b/autogl/module/nas/algorithm/rl.py @@ -250,7 +250,7 @@ class RL(BaseNAS): def __init__( self, num_epochs=5, - device="cuda", + device="auto", log_frequency=None, grad_clip=5.0, entropy_weight=0.0001, @@ -429,7 +429,7 @@ class GraphNasRL(BaseNAS): def __init__( self, - device="cuda", + device="auto", num_epochs=10, log_frequency=None, grad_clip=5.0, diff --git a/autogl/module/nas/space/base.py b/autogl/module/nas/space/base.py index 8ff7901..e122ccb 100644 --- a/autogl/module/nas/space/base.py +++ b/autogl/module/nas/space/base.py @@ -98,7 +98,7 @@ class BoxModel(BaseModel): _logger = get_logger("space model") - def __init__(self, space_model, device=torch.device("cuda")): + def __init__(self, space_model, device=torch.device("auto")): super().__init__(init=True) self.init = True self.space = [] @@ -247,7 +247,7 @@ class BaseSpace(nn.Module): ) return layer - def wrap(self, device="cuda"): + def wrap(self, device="auto"): """ Return a BoxModel which wrap self as a model Used to pass to trainer diff --git a/autogl/module/nas/space/gasso.py b/autogl/module/nas/space/gasso.py index 0397cb5..9d4db84 100644 --- a/autogl/module/nas/space/gasso.py +++ b/autogl/module/nas/space/gasso.py @@ -256,12 +256,17 @@ class GassoSpace(BaseSpace): def keep_prediction(self): self.prediction = self.current_pred + def to(self, *args, **kwargs): + super().to(args, kwargs) + device = next(self.parameters()).device + self.alphas_normal = self.alphas_normal.to(device) + def initialize_alphas(self): num_ops = len(self.ops) self.alphas_normal = [] for i in range(self.steps): - self.alphas_normal.append(Variable(1e-3 * torch.randn(num_ops).cuda(), requires_grad=True)) + self.alphas_normal.append(Variable(1e-3 * torch.randn(num_ops), requires_grad=True)) self._arch_parameters = [ self.alphas_normal From dc66c7b5342f693a50e9f979e9008ffd1c90782d Mon Sep 17 00:00:00 2001 From: "general502570@outlook.com" Date: Tue, 28 Dec 2021 21:49:22 +0800 Subject: [PATCH 10/14] add some tutorial --- autogl/module/nas/algorithm/gasso.py | 24 ++++++++++++++++++++ autogl/module/nas/estimator/one_shot.py | 20 +++++++++++++++- autogl/module/nas/estimator/train_scratch.py | 7 ++++++ autogl/module/nas/utils.py | 11 +++++++++ 4 files changed, 61 insertions(+), 1 deletion(-) diff --git a/autogl/module/nas/algorithm/gasso.py b/autogl/module/nas/algorithm/gasso.py index fc50ed6..7669301 100644 --- a/autogl/module/nas/algorithm/gasso.py +++ b/autogl/module/nas/algorithm/gasso.py @@ -25,6 +25,30 @@ _logger = logging.getLogger(__name__) @register_nas_algo("gasso") class Gasso(BaseNAS): + """ + DARTS trainer. + + Parameters + ---------- + num_epochs : int + Number of epochs planned for training. + warmup_epochs : int + Number of epochs planned for warming up. + workers : int + Workers for data loading. + model_lr : float + Learning rate to optimize the model. + model_wd : float + Weight decay to optimize the model. + arch_lr : float + Learning rate to optimize the architecture. + stru_lr : float + Learning rate to optimize the structure. + lamb : float + The parameter to control the influence of hidden feature smoothness + device : str or torch.device + The device of the whole process + """ def __init__( self, num_epochs=250, diff --git a/autogl/module/nas/estimator/one_shot.py b/autogl/module/nas/estimator/one_shot.py index 0081ee3..11b1aed 100644 --- a/autogl/module/nas/estimator/one_shot.py +++ b/autogl/module/nas/estimator/one_shot.py @@ -14,6 +14,13 @@ class OneShotEstimator(BaseEstimator): One shot estimator. Use model directly to get estimations. + + Parameters + ---------- + loss_f : str + The name of loss funciton in PyTorch + evaluation : list of Evaluation + The evaluation metrics in module/train/evaluation """ def __init__(self, loss_f="nll_loss", evaluation=[Acc()]): @@ -42,7 +49,18 @@ class OneShotEstimator_HardwareAware(OneShotEstimator): """ One shot hardware-aware estimator. - Use model directly to get estimations. + Use model directly to get estimations with some hardware-aware metrics. + + Parameters + ---------- + loss_f : str + The name of loss funciton in PyTorch + evaluation : list of Evaluation + The evaluation metrics in module/train/evaluation + hardware_evaluation : str or runable + The hardware-aware metrics. Can be 'parameter' or 'latency'. Or you can define a special metric by a runable function + hardware_metric_weight : float + The weight of hardware-aware metric, which will be a bias added to metrics """ def __init__( diff --git a/autogl/module/nas/estimator/train_scratch.py b/autogl/module/nas/estimator/train_scratch.py index 5677be7..84111c1 100644 --- a/autogl/module/nas/estimator/train_scratch.py +++ b/autogl/module/nas/estimator/train_scratch.py @@ -11,6 +11,13 @@ from autogl.module.train import NodeClassificationFullTrainer, Acc class TrainEstimator(BaseEstimator): """ An estimator which trans from scratch + + Parameters + ---------- + loss_f : str + The name of loss funciton in PyTorch + evaluation : list of Evaluation + The evaluation metrics in module/train/evaluation """ def __init__(self, loss_f="nll_loss", evaluation=[Acc()]): diff --git a/autogl/module/nas/utils.py b/autogl/module/nas/utils.py index 45f546f..a6ea9f0 100644 --- a/autogl/module/nas/utils.py +++ b/autogl/module/nas/utils.py @@ -101,6 +101,17 @@ class PathSamplingInputChoice(nn.Module): def get_hardware_aware_metric(model, hardware_metric): + """ + Get architectures' hardware-aware metrics + + Attributes + ---------- + model : BaseSpace + The architecture to be evaluated + hardware_metric : str + The name of hardware-aware metric. Can be 'parameter' or 'latency' + """ + if hardware_metric == 'parameter': return count_parameters(model) elif hardware_metric == 'latency': From 3854709e87483d95b7b4fb059e42982756abee8d Mon Sep 17 00:00:00 2001 From: "general502570@outlook.com" Date: Tue, 28 Dec 2021 22:29:25 +0800 Subject: [PATCH 11/14] fix sth --- autogl/module/nas/algorithm/enas.py | 1 - autogl/module/nas/algorithm/gasso.py | 2 +- autogl/module/nas/estimator/one_shot.py | 6 +++--- autogl/module/nas/space/base.py | 9 ++++----- autogl/module/nas/space/gasso.py | 4 ++-- autogl/module/nas/space/graph_nas.py | 2 +- autogl/module/nas/space/graph_nas_macro.py | 2 +- docs/docfile/tutorial/t_nas.rst | 4 ++-- 8 files changed, 14 insertions(+), 16 deletions(-) diff --git a/autogl/module/nas/algorithm/enas.py b/autogl/module/nas/algorithm/enas.py index d2c8dc3..122ee2d 100644 --- a/autogl/module/nas/algorithm/enas.py +++ b/autogl/module/nas/algorithm/enas.py @@ -101,7 +101,6 @@ class Enas(BaseNAS): def search(self, space: BaseSpace, dset, estimator): self.model = space self.dataset = dset # .to(self.device) - print(dir(dset)) self.estimator = estimator # replace choice self.nas_modules = [] diff --git a/autogl/module/nas/algorithm/gasso.py b/autogl/module/nas/algorithm/gasso.py index 7669301..1b8ac4f 100644 --- a/autogl/module/nas/algorithm/gasso.py +++ b/autogl/module/nas/algorithm/gasso.py @@ -26,7 +26,7 @@ _logger = logging.getLogger(__name__) @register_nas_algo("gasso") class Gasso(BaseNAS): """ - DARTS trainer. + GASSO trainer. Parameters ---------- diff --git a/autogl/module/nas/estimator/one_shot.py b/autogl/module/nas/estimator/one_shot.py index 11b1aed..c5dd593 100644 --- a/autogl/module/nas/estimator/one_shot.py +++ b/autogl/module/nas/estimator/one_shot.py @@ -14,7 +14,7 @@ class OneShotEstimator(BaseEstimator): One shot estimator. Use model directly to get estimations. - + Parameters ---------- loss_f : str @@ -37,8 +37,8 @@ class OneShotEstimator(BaseEstimator): y = label[mask] loss = getattr(F, self.loss_f)(pred, y) - #probs = F.softmax(pred, dim=1).detach().cpu().numpy() - probs = pred.detach().cpu().numpy() + probs = F.softmax(pred, dim=1).detach().cpu().numpy() + y = y.cpu() metrics = [eva.evaluate(probs, y) for eva in self.evaluation] return metrics, loss diff --git a/autogl/module/nas/space/base.py b/autogl/module/nas/space/base.py index 73bd39b..af32dcb 100644 --- a/autogl/module/nas/space/base.py +++ b/autogl/module/nas/space/base.py @@ -98,16 +98,15 @@ class BoxModel(BaseAutoModel): _logger = get_logger("space model") - def __init__(self, space_model, device=torch.device("auto")): + def __init__(self, space_model, device): super().__init__(None, None, device) self.init = True self.space = [] self.hyperparams = {} - self._model = space_model.to(device) + self._model = space_model self.num_features = self._model.input_dim self.num_classes = self._model.output_dim self.params = {"num_class": self.num_classes, "features_num": self.num_features} - self.device = device self.selection = None def _initialize(self): @@ -139,7 +138,6 @@ class BoxModel(BaseAutoModel): ret_self._model.instantiate() if ret_self.selection: apply_fixed_architecture(ret_self._model, ret_self.selection, verbose=False) - ret_self.to_device(self.device) return ret_self def __repr__(self) -> str: @@ -240,12 +238,13 @@ class BaseSpace(nn.Module): ) return layer - def wrap(self, device="auto"): + def wrap(self): """ 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` """ + device = next(self.parameters()).device return BoxModel(self, device) diff --git a/autogl/module/nas/space/gasso.py b/autogl/module/nas/space/gasso.py index 9d4db84..dc2be44 100644 --- a/autogl/module/nas/space/gasso.py +++ b/autogl/module/nas/space/gasso.py @@ -69,7 +69,7 @@ def gnn_map(gnn_name, in_dim, out_dim, concat=False, bias=True) -> Module: elif gnn_name == "zero": return ZeroConv(in_dim, out_dim, bias=bias) else: - print(gnn_name) + raise ValueError("No such GNN name") def Get_edges(adjs, ): edges = [] @@ -277,4 +277,4 @@ class GassoSpace(BaseSpace): def parse_model(self, selection, device) -> BaseModel: self.use_forward = False - return self.wrap(device) + return self.wrap() diff --git a/autogl/module/nas/space/graph_nas.py b/autogl/module/nas/space/graph_nas.py index 13dd559..fa4fca9 100644 --- a/autogl/module/nas/space/graph_nas.py +++ b/autogl/module/nas/space/graph_nas.py @@ -208,4 +208,4 @@ class GraphNasNodeClassificationSpace(BaseSpace): def parse_model(self, selection, device) -> BaseAutoModel: # return AutoGCN(self.input_dim, self.output_dim, device) - return self.wrap(device).fix(selection) + return self.wrap().fix(selection) diff --git a/autogl/module/nas/space/graph_nas_macro.py b/autogl/module/nas/space/graph_nas_macro.py index eab8ecc..6db7e52 100644 --- a/autogl/module/nas/space/graph_nas_macro.py +++ b/autogl/module/nas/space/graph_nas_macro.py @@ -183,7 +183,7 @@ class GraphNasMacroNodeClassificationSpace(BaseSpace): multi_label=False, batch_normal=False, layers=self.layer_number, - ).wrap(device) + ).wrap() return model diff --git a/docs/docfile/tutorial/t_nas.rst b/docs/docfile/tutorial/t_nas.rst index 6fbaad8..c0243b8 100644 --- a/docs/docfile/tutorial/t_nas.rst +++ b/docs/docfile/tutorial/t_nas.rst @@ -117,7 +117,7 @@ Here is an example. # For one-shot fashion, you can directly use following scheme in ``parse_model`` def parse_model(self, selection, device) -> BaseModel: - return self.wrap(device).fix(selection) + return self.wrap().fix(selection) Also, you can use the way which does not support one shot fashion. In this way, you can directly copy you model with few changes. @@ -153,7 +153,7 @@ But you can only use sample-based search strategy. # For non-one-shot fashion, you can directly return your model based on the choices # ``YourModel`` must inherit BaseSpace. def parse_model(self, selection, device) -> BaseModel: - model = YourModel(selection, self.input_dim, self.output_dim).wrap(device) + model = YourModel(selection, self.input_dim, self.output_dim).wrap() return model # YourModel can be defined as follows From cbd46cf66c351801eaf2b4290a5d7307cecc4a9b Mon Sep 17 00:00:00 2001 From: Frozenmad Date: Wed, 29 Dec 2021 10:27:39 +0800 Subject: [PATCH 12/14] fix bugs of import --- autogl/module/nas/algorithm/gasso.py | 2 +- autogl/module/nas/space/gasso.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/autogl/module/nas/algorithm/gasso.py b/autogl/module/nas/algorithm/gasso.py index 1b8ac4f..c3f5ba8 100644 --- a/autogl/module/nas/algorithm/gasso.py +++ b/autogl/module/nas/algorithm/gasso.py @@ -12,7 +12,7 @@ from .base import BaseNAS from ..estimator.base import BaseEstimator from ..space import BaseSpace from ..utils import replace_layer_choice, replace_input_choice -from ...model.base import BaseModel +from ...model.base import BaseAutoModel from torch.autograd import Variable import numpy as np diff --git a/autogl/module/nas/space/gasso.py b/autogl/module/nas/space/gasso.py index dc2be44..1248949 100644 --- a/autogl/module/nas/space/gasso.py +++ b/autogl/module/nas/space/gasso.py @@ -3,7 +3,7 @@ import typing as _typ from . import register_nas_space from .base import apply_fixed_architecture from .base import BaseSpace -from ...model import BaseModel +from ...model import BaseAutoModel from ....utils import get_logger from ..backend import * @@ -257,7 +257,7 @@ class GassoSpace(BaseSpace): self.prediction = self.current_pred def to(self, *args, **kwargs): - super().to(args, kwargs) + super().to(*args, **kwargs) device = next(self.parameters()).device self.alphas_normal = self.alphas_normal.to(device) @@ -275,6 +275,6 @@ class GassoSpace(BaseSpace): def arch_parameters(self): return self.alphas_normal - def parse_model(self, selection, device) -> BaseModel: + def parse_model(self, selection, device) -> BaseAutoModel: self.use_forward = False return self.wrap() From 68e5a21daad47d86c4ba77910f306f50e5ed6299 Mon Sep 17 00:00:00 2001 From: generall Date: Wed, 29 Dec 2021 15:42:35 +0800 Subject: [PATCH 13/14] fix gasso and add log in graphnas --- autogl/module/nas/algorithm/rl.py | 2 ++ autogl/module/nas/space/gasso.py | 9 +++++---- examples/graphnas.py | 2 +- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/autogl/module/nas/algorithm/rl.py b/autogl/module/nas/algorithm/rl.py index d927c26..2ba93a1 100644 --- a/autogl/module/nas/algorithm/rl.py +++ b/autogl/module/nas/algorithm/rl.py @@ -580,6 +580,8 @@ class GraphNasRL(BaseNAS): self.ctrl_optim.step() bar.set_postfix(acc=metric, max_acc=max(rewards)) + + LOGGER.info(f"epoch:{}, mean rewards:{}".format(epoch, sum(rewards) / len(rewards))) return sum(rewards) / len(rewards) def _resample(self): diff --git a/autogl/module/nas/space/gasso.py b/autogl/module/nas/space/gasso.py index 1248949..5d424f8 100644 --- a/autogl/module/nas/space/gasso.py +++ b/autogl/module/nas/space/gasso.py @@ -256,10 +256,11 @@ class GassoSpace(BaseSpace): def keep_prediction(self): self.prediction = self.current_pred - def to(self, *args, **kwargs): - super().to(*args, **kwargs) - device = next(self.parameters()).device - self.alphas_normal = self.alphas_normal.to(device) + '''def to(self, *args, **kwargs): + fin = super().to(*args, **kwargs) + device = next(fin.parameters()).device + fin.alphas_normal = [i.to(device) for i in self.alphas_normal] + return fin''' def initialize_alphas(self): num_ops = len(self.ops) diff --git a/examples/graphnas.py b/examples/graphnas.py index 2c49545..d12c5d0 100644 --- a/examples/graphnas.py +++ b/examples/graphnas.py @@ -12,7 +12,7 @@ if __name__ == '__main__': args = parser.parse_args() - dataset = build_dataset_from_name('cora') + dataset = build_dataset_from_name(args.dataset) label = dataset[0].nodes.data["y" if DependentBackend.is_pyg() else "label"][dataset[0].nodes.data["test_mask"]].cpu().numpy() solver = AutoNodeClassifier.from_config(args.config) solver.fit(dataset) From 4819f6c6ac76252dc48057d77455520f7d5f0cf5 Mon Sep 17 00:00:00 2001 From: wondergo2017 Date: Wed, 29 Dec 2021 17:25:02 +0800 Subject: [PATCH 14/14] add doc ; add test nclf --- autogl/module/nas/space/autoattend.py | 20 +++- test/nas/node_classification.py | 139 ++++++++++++++------------ 2 files changed, 95 insertions(+), 64 deletions(-) diff --git a/autogl/module/nas/space/autoattend.py b/autogl/module/nas/space/autoattend.py index 650726f..9441365 100644 --- a/autogl/module/nas/space/autoattend.py +++ b/autogl/module/nas/space/autoattend.py @@ -1,4 +1,4 @@ -# codes in this file are reproduced from https://github.com/GraphNAS/GraphNAS with some changes. +# codes in this file are reproduced from AutoAttend with some changes. from nni.nas.pytorch.mutables import Mutable import typing as _typ import torch @@ -24,10 +24,26 @@ OPS = [OPS1, OPS2] @register_nas_space("autoattend") class AutoAttendNodeClassificationSpace(BaseSpace): + """ + AutoAttend Search Space , please refer to http://proceedings.mlr.press/v139/guan21a.html for details. + The current implementation is nc (no context weight sharing), + we will in future add other types of partial weight sharing proposed in the paper. + + Parameters + ---------- + ops_type : int + 0 or 1 , choosing from two sets of ops with index ops_type + gnn_ops : list of str + op names for searching, which descripts the compostion of operation pool + act_op : str + determine used activation function + agg_ops : list of str + agg op names for searching. Only ['add','attn'] are options, as mentioned in the paper. + """ def __init__( self, hidden_dim: _typ.Optional[int] = 64, - layer_number: _typ.Optional[int] = 5, + layer_number: _typ.Optional[int] = 2, dropout: _typ.Optional[float] = 0.9, input_dim: _typ.Optional[int] = None, output_dim: _typ.Optional[int] = None, diff --git a/test/nas/node_classification.py b/test/nas/node_classification.py index a60bb44..819b4d8 100644 --- a/test/nas/node_classification.py +++ b/test/nas/node_classification.py @@ -82,14 +82,13 @@ if __name__ == "__main__": di = bk_feat(data).shape[1] do = len(np.unique(bk_label(data))) - # print("Random search + singlepath ") - # space = SinglePathNodeClassificationSpace().cuda() - # space.instantiate(input_dim=di, output_dim=do) - # esti = OneShotEstimator() - # algo = RandomSearch(num_epochs=100) - # model = algo.search(space, dataset, esti) - # test_model(model, data, True) - + print("Random search + singlepath ") + space = SinglePathNodeClassificationSpace().cuda() + space.instantiate(input_dim=di, output_dim=do) + esti = OneShotEstimator() + algo = RandomSearch(num_epochs=100) + model = algo.search(space, dataset, esti) + test_model(model, data, True) print("Random search + AutoAttend ") space = AutoAttendNodeClassificationSpace().cuda() @@ -100,58 +99,74 @@ if __name__ == "__main__": print(model) test_model(model, data, True) - # print("Random search + graphnas ") - # space = GraphNasNodeClassificationSpace().cuda() - # space.instantiate(input_dim=di, output_dim=do) - # esti = OneShotEstimator() - # algo = RandomSearch(num_epochs=10) - # model = algo.search(space, dataset, esti) - # test_model(model, data, True) - - # print("rl + graphnas ") - # space = GraphNasNodeClassificationSpace().cuda() - # space.instantiate(input_dim=di, output_dim=do) - # esti = OneShotEstimator() - # algo = RL(num_epochs=10) - # model = algo.search(space, dataset, esti) - # test_model(model, data, True) - - # print("graphnasrl + graphnas ") - # space = GraphNasNodeClassificationSpace().cuda() - # space.instantiate(input_dim=di, output_dim=do) - # esti = OneShotEstimator() - # algo = GraphNasRL(num_epochs=10) - # model = algo.search(space, dataset, esti) - # test_model(model, data, True) - - # print("darts + graphnas ") - # space = GraphNasNodeClassificationSpace(con_ops=['concat']).cuda() - # space.instantiate(input_dim=di, output_dim=do) - # esti = OneShotEstimator() - # algo = Darts(num_epochs=10) - # model = algo.search(space, dataset, esti) - # test_model(model, data, True) + print("rl + AutoAttend ") + space = AutoAttendNodeClassificationSpace().cuda() + space.instantiate(input_dim=di, output_dim=do) + esti = OneShotEstimator() + algo = RL(num_epochs=10) + model = algo.search(space, dataset, esti) + test_model(model, data, True) + + print("darts + graphnas ") + space = AutoAttendNodeClassificationSpace().cuda() + space.instantiate(input_dim=di, output_dim=do) + esti = OneShotEstimator() + algo = Darts(num_epochs=10) + model = algo.search(space, dataset, esti) + test_model(model, data, True) + + print("Random search + graphnas ") + space = GraphNasNodeClassificationSpace().cuda() + space.instantiate(input_dim=di, output_dim=do) + esti = OneShotEstimator() + algo = RandomSearch(num_epochs=10) + model = algo.search(space, dataset, esti) + test_model(model, data, True) + + print("rl + graphnas ") + space = GraphNasNodeClassificationSpace().cuda() + space.instantiate(input_dim=di, output_dim=do) + esti = OneShotEstimator() + algo = RL(num_epochs=10) + model = algo.search(space, dataset, esti) + test_model(model, data, True) + + print("graphnasrl + graphnas ") + space = GraphNasNodeClassificationSpace().cuda() + space.instantiate(input_dim=di, output_dim=do) + esti = OneShotEstimator() + algo = GraphNasRL(num_epochs=10) + model = algo.search(space, dataset, esti) + test_model(model, data, True) + + print("darts + graphnas ") + space = GraphNasNodeClassificationSpace(con_ops=['concat']).cuda() + space.instantiate(input_dim=di, output_dim=do) + esti = OneShotEstimator() + algo = Darts(num_epochs=10) + model = algo.search(space, dataset, esti) + test_model(model, data, True) - # print("darts + singlepath ") - # space = SinglePathNodeClassificationSpace().cuda() - # space.instantiate(input_dim=di, output_dim=do) - # esti = OneShotEstimator() - # algo = Darts(num_epochs=10) - # model = algo.search(space, dataset, esti) - # test_model(model, data, True) - - # print("Random search + graphnas macro") - # space = GraphNasMacroNodeClassificationSpace().cuda() - # space.instantiate(input_dim=di, output_dim=do) - # esti = OneShotEstimator() - # algo = RandomSearch(num_epochs=10) - # model = algo.search(space, dataset, esti) - # test_model(model, data, True) - - # print("RL + graphnas macro") - # space = GraphNasMacroNodeClassificationSpace().cuda() - # space.instantiate(input_dim=di, output_dim=do) - # esti = OneShotEstimator() - # algo = RL(num_epochs=10) - # model = algo.search(space, dataset, esti) - # test_model(model, data, True) + print("darts + singlepath ") + space = SinglePathNodeClassificationSpace().cuda() + space.instantiate(input_dim=di, output_dim=do) + esti = OneShotEstimator() + algo = Darts(num_epochs=10) + model = algo.search(space, dataset, esti) + test_model(model, data, True) + + print("Random search + graphnas macro") + space = GraphNasMacroNodeClassificationSpace().cuda() + space.instantiate(input_dim=di, output_dim=do) + esti = OneShotEstimator() + algo = RandomSearch(num_epochs=10) + model = algo.search(space, dataset, esti) + test_model(model, data, True) + + print("RL + graphnas macro") + space = GraphNasMacroNodeClassificationSpace().cuda() + space.instantiate(input_dim=di, output_dim=do) + esti = OneShotEstimator() + algo = RL(num_epochs=10) + model = algo.search(space, dataset, esti) + test_model(model, data, True)