为机器学习打开 ROOT NTuple 的最快、最节省内存的方法是什么？

Question

我正在使用 scikit learn 建立一个机器学习项目。输入数据是平面 ROOT NTuples。

过去我一直使用 root_numpy 将 NTuples 转换为保存在 h5 文件中的 pandas.DataFrame 。

我想知道我是否可以使用 uproot 来
a) 完全跳过 h5 转换？
b) 不使用与从 h5 加载到 DataFrame 中一样多的内存？

我天真的第一次尝试看起来像这样：

'''
Runs preselection, keeps only desired variables in DataFrame
'''
def dropAndKeep(df, dropVariables=None, keepVariables=None, presel=None, inplace=True):

    if ((presel is not None) and (not callable(presel))):
        print("Please either provide a function to 'presel' or leave blank")
        raise ValueError

    if callable(presel):
        if not(inplace):
            df = df.drop(df[~presel(df)].index, inplace=False)
        else:
            df.drop(df[~presel(df)].index, inplace=True)

    if keepVariables is not None:
        dropThese = list( set(df.columns) - set(keepVariables) )
        return df.drop(columns=dropThese, inplace=inplace)

    if dropVariables is not None:
        return df.drop(columns=dropVariables, inplace=inplace)

'''
Loads a TTree from ROOT file into a DataFrame 
'''
def load_root(inFile, key, dropVariables=None, keepVariables=None, presel=None):
    df = uproot.open(inFile)[key].pandas.df()
    dropAndKeep(df, dropVariables, keepVariables, presel=presel, inplace=True)
    return df


inFile = "path/to/file.root"
key = "ntuple"
df = load_root(inFile, key)

这需要很长时间。有更好的方法吗？

Answer 1

请注意，每次调用uproot.open(...)和file [key]加载 TFile 和 TTree 元数据是 uproot 中最慢的部分。如果多次调用此函数，请尝试保留 TFile 和/或 TTree 对象并重新使用它们。

\n\n

另外，看起来你的dropAndKeep函数只是删除行（事件），但如果我读错了并且它正在执行列（分支），那么使用branchesuproot 的数组读取函数的参数仅发送你想要的分支。由于 ROOT 文件中的数据按列排列，因此您无法避免读取不需要的事件\xe2\x80\x94，您必须在事后删除它们（在任何框架中）。

\n\n

接下来，请注意，对于过滤事件等简单操作，Pandas 比 NumPy 慢得多。如果您想加快速度，请使用TTree.arrays，而不是获取数组TTree.pandas.df，为您的选择构造一个布尔值的 NumPy 数组，并将其应用于字典中的每个数组TTree.arrays。然后，您可以使用 Pandas 的 DataFrame 构造函数将所有这些放入 DataFrame 中（如果您确实需要 Pandas）。

\n\n

确实，您不需要通过 HDF5，也不需要通过 Pandas。您的机器学习框架（TensorFlow？Torch？）几乎肯定有一个接口，可以接受零拷贝（或到 GPU 的一次拷贝）的 NumPy 数组。强调 HDF5 或 Pandas 的教程这样做是因为对于大多数用户（非 HEP）来说，这些是最方便的界面。他们的数据可能已经在 HDF5 或 Pandas 中；我们的数据可能在ROOT中。

\n\n

如果您的机器学习将在 GPU 上进行，也许您也想在 GPU 上进行事件选择。CuPy是 NumPy 克隆，完全在 GPU 上分配和操作，并且您的 TensorFlow/Torch 张量可能具有与 CuPy 数组的零拷贝接口。原则上，如果使用 CuPy 数组作为asarray 解释的目标，uproot 应该能够直接从 ROOT 文件写入 CuPy 数组。不过我还没有尝试过。

\n\n

如果您可以控制要处理的 ROOT 文件，请尝试使它们的篮子变大（增加刷新大小）并使它们的数据结构简单（例如纯数字或数字数组/向量，不要更深）。也许最重要的是，使用像 lz4 这样的轻量级压缩，而不是重量级的 Luke lzma。

\n\n

Uproot 可以并行读取篮子，但只有在需要执行大量非 Python 计算（例如解压缩 lzma）时，这才被证明是有用的。

\n\n

如果您要一遍又一遍地读取这些数组，您可能需要使用numpy.save写入中间文件，这本质上只是磁盘上的原始字节。这意味着读回时不需要反序列化，这与解码 ROOT 或 HDF5 文件所需的工作相反。因为它的格式非常简单，您甚至可以使用numpy.memmap读回它读回它，它会在从磁盘延迟加载数据时查看操作系统的页面缓存，甚至删除字节的显式副本。

\n\n

并非所有这些技巧都同样有用。我试图把最重要的放在第一位，但在进行大的代码重写之前进行实验，这可能不会产生太大的影响。有些技巧不能与其他技巧结合使用，例如 CuPy 和 memmap（memmap 总是延迟加载到主内存中，而不是 GPU 内存中）。但某些组合可能会取得成果。

\n\n

祝你好运！

\n