python - 打开用于机器学习的 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] 使用纯 Python 加载 TFile 和 TTree 元数据——uproot 中最慢的部分.如果您多次调用此方法，请尝试保留 TFile 和/或 TTree 对象并重新使用它们。

此外，您的 dropAndKeep 函数看起来只是删除行(事件)，但如果我读错了它正在处理列(分支)，那么请使用 branches uproot 的数组读取函数的参数只发送你想要的分支。由于 ROOT 文件中的数据按列排列，您无法避免读取不需要的事件 - 您必须事后删除它们(在任何框架中)。

接下来，请注意，对于过滤事件等简单操作，Pandas 比 NumPy 慢得多。如果您想加快速度，请使用 TTree.arrays 而不是 TTree.pandas.df 获取数组，构建一个 NumPy bool 数组供您选择，然后应用它到 TTree.arrays 返回的字典中的每个数组。然后，您可以使用 Pandas 的 DataFrame 构造函数将所有这些放入 DataFrame 中(如果您真的需要 Pandas)。

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

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

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

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

如果您打算一遍又一遍地读取这些数组，您可能希望使用 numpy.save 编写中间文件，这实际上只是磁盘上的原始字节。这意味着在回读时没有反序列化，而不是解码 ROOT 或 HDF5 文件所需的工作。因为它是一种如此简单的格式，你甚至可以用 numpy.memmap 读回它。 ，当它从磁盘延迟加载数据时，它会窥视操作系统的页面缓存，甚至删除字节的显式副本。

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

祝你好运!