Google - 快还是慢？预测AI模型运行时竞赛第13名解决方案

这是我参加"快还是慢？预测AI模型运行时"竞赛的解决方案。希望能对你有所帮助！

我的方法核心原则是从简单方案开始并通过多次迭代改进，同时在硬件限制下工作（16 GB内存，4 GB显存）。

背景部分

业务背景: https://www.kaggle.com/competitions/predict-ai-model-runtime/overview
数据背景: https://www.kaggle.com/competitions/predict-ai-model-runtime/data

方法概述

布局模型

核心思想是从图和节点中提取特征，并用这些信息训练一个多层感知机。对于计算图的每个可配置节点，我们提取了一些节点属性、使用的布局，以及来自其"父节点"和"兄弟节点"的属性。"父节点"是产生节点输入的节点。如果两个节点共享一个父节点，则它们是"兄弟节点"。例如，在下图中，节点1和2是父节点，节点3、5和6是兄弟节点。

对于图中的每个可配置节点，我们都有一个特征列表，这些特征通过3个全连接层处理，层间有dropout层。之后，我们对"节点"维度进行平均，因此所有节点信息被表示为一个向量。有两个额外输入被连接到这个向量：一个"图描述"和"子集信息"。"图描述"是一个向量，包含图中每种类型节点的数量（归一化为总和为1），以及一个额外值表示图的节点数（取对数），让模型了解图的大小。"子集信息"是一个向量，表示图来自"xla:default"、"xla:random"、"nlp:default"还是"nlp:random"子集，为此我们使用了keras的"嵌入"层。

这个新向量再经过3个额外的全连接层处理（这次没有dropout）。训练期间使用成对Hinge损失作为目标函数。为此，每个大小为128的训练批次包含16个不同图的示例，每个图有8个配置示例。每个批次通过随机选择子集构建，平均而言批次是子集平衡的。最终提交结果是3次独立训练的集成。

分块模型

分块模型是布局模型的简化版本。每个配置用一个向量描述，包含"config_feat"信息和"图描述符"（与布局模型相同）。该模型是一个多层感知机，有3个全连接层，在第一层后有一个dropout层。损失函数和批次结构与布局模型相同，但批次更大（600）且每图配置更多（20）。

验证

使用数据集中的"valid"文件夹示例作为验证集。每10000次训练迭代进行一次验证，并计算该集合上的竞赛指标。对于布局模型，计算四个子集的指标后取平均值。如果验证指标在5次验证后没有改善，则停止训练。

提交详情

分块模型详情

"分块问题"相对"布局问题"较为简单，因此我们没有投入太多精力改进此模型。每图配置数限制为160，我们使用指数分布采样，给运行时较低的配置更高概率，因为这里的挑战是找到最快配置，而非对所有配置排序。

节点特征（布局模型）

从"node_feat"矩阵的所有可用特征中，我们选择了认为最重要的特征，这有助于保持低内存需求。这些特征包括：

shape_dimensions (21-26)
reshape/broadcast dimensions (31-36)
convolution_dim_numbers_input_spatial_dims (95-98)
convolution_dim_numbers_kernel_spatial_dims (101-104)
layout_minor_to_major (134-139)

括号中的值对应"node_feat"的选定索引。我们还提供了节点布局信息和操作码（通过keras嵌入层编码为向量）。另一个重要点是，根据布局信息，我们还向网络提供了形状的重排版本（除原始版本外）。

对于兄弟节点，我们获取每个兄弟节点的输出形状、布局，以及一个布尔值来比较节点布局与兄弟节点布局是否相同。正如竞赛概述中提到的，如果两个兄弟节点的布局不同，则需要额外的复制操作，这促使我们创建了这个变量。

对于父节点，我们保留它们的输出形状和物理布局。我们还从父节点和兄弟节点获取操作码，并通过keras嵌入层将其表示为向量。

保持训练稳定（布局模型）

为便于训练，所有取值范围跨越多个数量级的特征（如张量形状）都经过对数处理，以避免非常大的输入值。特征使用均值/标准差归一化，在标准差估计中进行一些截断，避免除以极小的值。

学习率使用余弦衰减调度。在10000次迭代的线性预热后，余弦衰减在250k次迭代中将参数降至原始值的5%。使用Adam优化器，clipnorm值为1.0，避免权重更新过大。

许多配置具有相同的布局。它们全部被替换为该布局的一个实例，运行时替换为运行时的平均值。

效果不佳的尝试

我们在List MLE损失上遇到了一些不稳定问题，因此选择使用成对Hinge损失。问题是经过多次迭代后，模型（或损失）中突然出现NaN值，摧毁了所有模型权重。

在内存预算内进行训练

由于我们的内存预算有限（16 GB内存，4 GB显存），必须非常小心不要同时加载过多数据，并保持模型大小保守。首要重要的是处理所有npz文件，将必要信息以tfrecords格式保存到tensorflow中（启用文件压缩）。训练期间，从磁盘读取这些文件，尝试给模型提供来自多个不同图的样本，而非一次只看一个图。

每图配置数上限为7500，给网络的可配置节点数上限为1000。由于给网络提供多个不同图，且每个图的可配置节点数不同，需要填充和掩码张量，使所有样本在"节点"维度上长度相同。如果增加使用的节点数，内存负担会很重，因此选择1000是考虑了这个限制。我还限制了每个节点的父节点数（2）和兄弟节点数（3）。