神经架构搜索

我们支持在不同搜索空间中应用不同的神经架构搜索算法。 神经架构搜索通常由三个模块构成: 搜索空间,搜索策略和评估策略。

搜索空间描述了所有可能被搜索的架构。空间主要由两部分组成,即操作(例如 GCNconv、GATconv)和输入-输出关系。 大的搜索空间可能有更好的最优架构,但需要更多的努力去探索。 人类的知识可以帮助设计一个合理的搜索空间,减少搜索策略的努力。

搜索策略控制了如何探索搜索空间, 它涵盖了经典的探索-利用权衡: 一方面,希望尽快找到性能良好的架构, 另一方面,应该避免过早地收敛到次优架构的区域。

评估策略在探索某一架构时给出其性能。 最简单的选项是在数据上对架构进行标准训练和验证。 由于在整个搜索过程中需要评估许多架构,评估策略需要非常高效以节约计算资源。

为了更灵活,我们将NAS过程模块化为三部分:搜索算法,搜索空间和评估器。 不同部分的不同模型可以在某些特定的约束下进行组合。 如果你想设计你自己的NAS过程,你可以根据你的需求改变其中任何一部分。

用法

你可以直接引入特定的空间、算法、估计器来为特定的数据集搜索图神经网络。下面是一个例子:

from autogllight.nas.space import GraphNasNodeClassificationSpace
from autogllight.nas.algorithm import GraphNasRL
from autogllight.nas.estimator import OneShotEstimator
from torch_geometric.datasets import Planetoid
from os import path as osp
import torch_geometric.transforms as T

dataname = "cora"
dataset = Planetoid(
    osp.expanduser("~/.cache-autogl"), dataname, transform=T.NormalizeFeatures()
)
data = dataset[0]
label = data.y
input_dim = data.x.shape[-1]
num_classes = len(np.unique(label.numpy()))

space = GraphNasNodeClassificationSpace(input_dim=input_dim, output_dim=num_classes)
space.instantiate()
algo = GraphNasRL(num_epochs=2, ctrl_steps_aggregate=2, weight_share=False)
estimator = OneShotEstimator()
algo.search(space, dataset, estimator)

以上代码在搜索空间 GraphNasNodeClassificationSpace 使用 GraphNasRL 搜索算法进行搜索。

搜索空间

搜索空间需要继承BaseSpace。 定义搜索空间主要有两种方式,一种可以以one-shot方式执行,另一种则不能。 目前,我们支持以下搜索空间: SinglePathNodeClassificationSpace, GassoSpace, GraphNasNodeClassificationSpace, GraphNasMacroNodeClassificationSpace, AutoAttendNodeClassificationSpace。

你也可以定义你自己的nas搜索空间。你应该重写``build_graph``函数来构建超网络。这里有一个例子。

from autogllight.space.base import BaseSpace

# For example, create an NAS search space by yourself
class SinglePathNodeClassificationSpace(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 = ["gcn", "gat_8"],
):
    super().__init__()
    self.layer_number = layer_number
    self.hidden_dim = hidden_dim
    self.input_dim = input_dim
    self.output_dim = output_dim
    self.ops = ops
    self.dropout = dropout

def build_graph(self):
    for layer in range(self.layer_number):
        key = f"op_{layer}"
        in_dim = self.input_dim if layer == 0 else self.hidden_dim
        out_dim = (self.output_dim if layer == self.layer_number - 1 else self.hidden_dim)
        op_candidates = [
            op(in_dim, out_dim)
            if isinstance(op, type)
            else gnn_map(op, in_dim, out_dim)
            for op in self.ops
        ]
        self.setLayerChoice(layer, op_candidates, key=key)

def forward(self, data):
    x = BK.feat(data)
    for layer in range(self.layer_number):
        op = getattr(self, f"op_{layer}")
        x = BK.gconv(op, data, x)
        if layer != self.layer_number - 1:
            x = F.leaky_relu(x)
            x = F.dropout(x, p=self.dropout, training=self.training)
    return F.log_softmax(x, dim=1)

性能评估

性能估计器用于估计架构的性能。目前我们支持以下估算器:

您也可以编写自己的估算器。这是一个在没有训练的情况下估计架构(用于one-shot空间)的示例。

# For example, create an NAS estimator by yourself
from autogllight.nas.estimator.base import BaseEstimator
class YourOneShotEstimator(BaseEstimator):
    # The only thing you should do is defining ``infer`` function
    def infer(self, model: BaseSpace, dataset, mask="train"):
        device = next(model.parameters()).device
        dset = dataset[0].to(device)
        # Forward the architecture
        pred = model(dset)[getattr(dset, f"{mask}_mask")]
        y = dset.y[getattr(dset, f'{mask}_mask')]
        # Use default loss function and metrics to evaluate the architecture
        loss = getattr(F, self.loss_f)(pred, y)
        probs = F.softmax(pred, dim = 1)
        metrics = [eva.evaluate(probs, y) for eva in self.evaluation]
        return metrics, loss

搜索空间

空间策略定义了如何寻找架构。我们目前支持以下搜索策略:RandomSearch、Darts、RL、GraphNasRL、Enas、Spos、GRNA、Gasso。

不共享权重的基于样本的策略比具有权重共享的策略更简单。我们以DFS作为示例来展示如何定义自己的策略。

from autogllight.nas.algorithm.base import BaseNAS
class RandomSearch(BaseNAS):
    # Get the number of samples at initialization
    def __init__(self, n_sample):
        super().__init__()
        self.n_sample = n_sample

    # The key process in NAS algorithm, search for an architecture given space, dataset and estimator
    def search(self, space: BaseSpace, dset, estimator):
        self.estimator=estimator
        self.dataset=dset
        self.space=space

        self.nas_modules = []
        k2o = get_module_order(self.space)
        # collect all mutables in the space
        replace_layer_choice(self.space, PathSamplingLayerChoice, self.nas_modules)
        replace_input_choice(self.space, PathSamplingInputChoice, self.nas_modules)
        # sort all mutables with given orders
        self.nas_modules = sort_replaced_module(k2o, self.nas_modules)
        # get a dict cantaining all chioces
        selection_range={}
        for k,v in self.nas_modules:
            selection_range[k]=len(v)
        self.selection_dict=selection_range

        arch_perfs=[]
        # define DFS process
        self.selection = {}
        last_k = list(self.selection_dict.keys())[-1]
        def dfs():
            for k,v in self.selection_dict.items():
                if not k in self.selection:
                    for i in range(v):
                        self.selection[k] = i
                        if k == last_k:
                            # evaluate an architecture
                            self.arch=space.parse_model(self.selection,self.device)
                            metric,loss=self._infer(mask='val')
                            arch_perfs.append([metric, self.selection.copy()])
                        else:
                            dfs()
                    del self.selection[k]
                    break
        dfs()

        # get the architecture with the best performance
        selection=arch_perfs[np.argmax([x[0] for x in arch_perfs])][1]
        arch=space.parse_model(selection,self.device)
        return arch

不同的搜索策略应与不同的搜索空间和估算器结合使用。大多数搜索空间、搜索策略和估算器是兼容的。