GNN学习笔记|AEGNN

参考文献AEGNN: Asynchronous Event-based Graph Neural Networks

先load数据集，再在setup中调整数据集的输入方式，最后训练模型。

首先是第一行代码

    data_module = aegnn.datasets.NCars(batch_size=1, shuffle=False)

这行代码实例化了data_module为Ncars类，并设置为每次处理一个样本（batch_size=1），并且不打乱数据顺序（shuffle=False），图像尺寸被设置为 (120, 100)，并保留了一系列的超参数，包括：

• r: 用于创建半径图时的半径值。
• d_max: 最大邻居数。
• n_samples: 从事件点云中采样的点的数量。
• sampling: 是否启用采样。

然后执行下一句代码

data_module.setup()

在setup中加载训练数据集和验证数据集，并对数据集的完整性进行检查。

最后执行

run_experiments(data_module, arguments, experiments=event_counts, num_trials=100,
                    device=torch.device(arguments.device), log_flops=True, log_runtime=True)

def run_experiments(dm, args, experiments: List[int], num_trials: int, device: torch.device, **model_kwargs
                    ) -> pd.DataFrame:
    results_df = pd.DataFrame()
    output_file = os.path.join(os.path.dirname(os.path.realpath(__file__)), "..", "..", "aegnn_results", "flops.pkl")
    os.makedirs(os.path.dirname(output_file), exist_ok=True)

    runs = list(itertools.product(experiments, list(range(num_trials))))
    for num_events, exp_id in tqdm(runs):
        model = create_and_run_model(dm, num_events, index=exp_id, args=args, device=device, **model_kwargs)

        # Get the logged flops and timings, both layer-wise and in total.
        results_flops = get_log_values(model, attr="asy_flops_log", log_key="flops", num_events=num_events)
        results_runtime = get_log_values(model, attr="asy_runtime_log", log_key="runtime", num_events=num_events)
        results_df = results_df.append(results_flops + results_runtime, ignore_index=True)
        results_df.to_pickle(output_file)

        # Fully reset run to ensure independence between subsequent experiments.
        del model  # fully delete model
        torch.cuda.empty_cache()  # clear memory

    print(f"Results are logged in {output_file}")
    return results_df

先初始化DataFrame，用于储存实验结果。

再设置输出文件路径，如果输出目录不存在，则创建一个。

再生成所有实验组合，其中 experiments 是一个事件数量的列表，num_trials 指定了每个事件数量需要进行多少次试验。

如果

experiments = ['A', 'B', 'C']  # 三个不同的实验配置
num_trials = 2  # 每个实验配置重复两次

则输出为

[('A', 0), ('A', 1), ('B', 0), ('B', 1), ('C', 0), ('C', 1)]

然后进行训练，使用tqdm显示实验进度，并根据给定的事件数量 (num_events) 和试验 ID (exp_id) 创建模型。

model = create_and_run_model(dm, num_events, index=exp_id, args=args, device=device, **model_kwargs)

def create_and_run_model(dm, num_events: int, index: int, device: torch.device, args: argparse.Namespace, **kwargs):
    edge_attr = torch_geometric.transforms.Cartesian(cat=False, max_value=10.0)
    dataset = dm.train_dataset
    assert dm.shuffle is False  # ensuring same samples over experiments

    # Sample initial data of certain length from dataset sample. Sample num_events samples from one
    # dataset, and create the subsequent event as the one to be added.
    sample = dataset[index % len(dataset)]
    sample.pos = sample.pos[:, :2]
    events_initial = sample_initial_data(sample, num_events, args.radius, edge_attr, args.max_num_neighbors)

    index_new = min(num_events, sample.num_nodes - 1)
    x_new = sample.x[index_new, :].view(1, -1)
    pos_new = sample.pos[index_new, :2].view(1, -1)
    event_new = Data(x=x_new, pos=pos_new, batch=torch.zeros(1, dtype=torch.long))

    # Initialize model and make it asynchronous (recognition model, so num_outputs = num_classes of input dataset).
    input_shape = torch.tensor([*dm.dims, events_initial.pos.shape[-1]], device=device)
    model = aegnn.models.networks.GraphRes(dm.name, input_shape, dm.num_classes, pooling_size=args.pooling_size)
    model.to(device)
    model = aegnn.asyncronous.make_model_asynchronous(model, args.radius, list(dm.dims), edge_attr, **kwargs)

    # Run experiment, i.e. initialize the asynchronous graph and iteratively add events to it.
    _ = model.forward(events_initial.to(device))  # initialization
    _ = model.forward(event_new.to(device))
    del events_initial, event_new
    return model

先定义一个边属性变换edge_attr，用torch_geometric.transforms.Cartesian 类来计算节点之间的相对位置信息。

举例来说

edge_index = torch.tensor([[0, 1, 1, 2],
                           [1, 0, 2, 1]], dtype=torch.long)
pos = torch.tensor([[-1, -1], [-1, 1], [1, 1]], dtype=torch.float)

• edge_index: 定义了图中的边连接关系。例如，边 (0, 1) 表示节点 0 和节点 1 之间有一条边。
• pos: 定义了每个节点的位置。例如，节点 0 的位置是 (-1, -1)，节点 1 的位置是 (-1, 1)，节点 2 的位置是 (1, 1)。
• edge_attr: 计算了每条边的相对位置。例如：
  • 边 (0, 1) 的相对位置是从节点 0 到节点 1 的位置差，即 (0 - (-1), 1 - (-1)) = (0, 2)。
  • 边 (1, 0) 的相对位置是从节点 1 到节点 0 的位置差，即 ((-1) - (-1), (-1) - 1) = (0, -2)。
  • 边 (1, 2) 的相对位置是从节点 1 到节点 2 的位置差，即 (1 - (-1), 1 - 1) = (2, 0)。
  • 边 (2, 1) 的相对位置是从节点 2 到节点 1 的位置差，即 ((-1) - 1, 1 - 1) = (-2, 0)。

然后从dm中获得训练数据集，并设置dm.shuffle为false，确保在不同实验中样本顺序相同。

从数据集中获取指定索引样本，并仅保留前两个维度的位置信息。从样本中选择前 num_events 个事件作为初始数据，并创建下一个事件作为要添加的数据。

这里的每个事件应该是指每个像素点在某个时刻的极性吧。

def sample_initial_data(sample, num_events: int, radius: float, edge_attr, max_num_neighbors: int):
    data = Data(x=sample.x[:num_events], pos=sample.pos[:num_events])
    data.batch = torch.zeros(data.num_nodes, device=data.x.device)
    data.edge_index = torch_geometric.nn.radius_graph(data.pos, r=radius, max_num_neighbors=max_num_neighbors).long()
    data.edge_attr = edge_attr(data).edge_attr

    edge_counts_avg = data.edge_index.shape[1] / num_events
    logging.debug(f"Average edge counts in initial data = {edge_counts_avg}")
    return data

他这里设置的batch_size是1，应该是一次导入一张图的意思吧，但是一张图的大小是120x100，event_counts是25000，不太理解。

data包含像素点位置以及极性的信息。

data.batch = torch.zeros(data.num_nodes, device=data.x.device)

然后初始化批次信息，这里的批次信息通常用于指示哪些节点属于同一个图。在这个例子中，所有节点都属于同一个批次，因此批次信息是一个全零的张量。有点明白了，应该是多张图的时候，这个25000有用。

再根据节点距离确定边的存在，并且限制每个节点的最大邻居数量为 max_num_neighbors。使用提供的 edge_attr 函数或变换对象来计算边的属性。这通常包括节点之间的相对位置或其他相关信息。计算每个节点的平均边数量，并记录日志信息。返回包含节点特征、位置、边索引和边属性的 Data 对象。

然后再确定新的事件索引，并提取新的事件特征和事件位置，最后创建一个Data对象。

再定义模型的输入形状，初始化模型。

模型等我跑一遍再来分析

然后再使模型异步