如何将非锯齿Awkward数组转换为PyTorch张量（v1.10.3）

本人使用Awkward版本1.10.3。现有一组存储在Awkward数组中的非锯齿数据，需要传入PyTorch前馈模型。由于PyTorch不原生支持Awkward数组，计划先转为Torch或Numpy数组。

示例数据及预期转换结果如下：

import awkward as ak
import numpy as np
import torch

arr = ak.Array({"MET_pt" : [0.0, 100.0, 20.0, 30.0, 4.0],
                 "MET_phi" : [0, 0.1, 0.2, 0.3, 0.4],
                 "class" : [0, 1, 0, 1, 0]})

# 输入特征
x = arr[['MET_pt', 'MET_phi']]
# 类别标签
y = arr['class']

# 预期转换后的张量
x_torch = torch.tensor([[0, 0], [100, 0.1], [20, 0.2], [30, 0.3], [4, 0.4]])
y_torch = torch.tensor([0, 1, 0, 1, 0])

转换特征数组x时遇到以下问题：

• 直接调用torch.tensor(x)报错：TypeError: object of type 'Record' has no len()
• 先通过ak.to_numpy(x)转为Numpy数组得到结构化数组，无法转为Torch张量，报错：TypeError: can't convert np.ndarray of type numpy.void...

需求：将x转为二维Numpy数组，或直接转为目标Torch张量。

方法1：结构化数组转二维数组后转张量 #

先将Awkward数组转为Numpy结构化数组，再通过view转换为普通二维数组，最后转Torch张量：

# 转Numpy结构化数组
x_np = ak.to_numpy(x)
# 转换为二维数组（确保dtype统一，这里均为float）
x_np_2d = x_np.view(np.float64).reshape(x_np.shape[0], -1)
# 转Torch张量
x_torch = torch.from_numpy(x_np_2d)

方法2：用ak.unflatten构造二维结构后直接转张量 #

利用ak.unflatten将多字段合并为二维Awkward数组，Awkward 1.10+支持直接转为Torch张量：

# 合并字段为二维结构
x_2d = ak.unflatten(ak.flatten(x), 2)
# 直接转Torch张量
x_torch = torch.tensor(x_2d)

方法3：手动提取字段拼接 #

逐个提取特征字段的Numpy数组，拼接为二维数组后转张量：

# 提取单个字段的Numpy数组
pt_np = ak.to_numpy(x['MET_pt'])
phi_np = ak.to_numpy(x['MET_phi'])
# 拼接为二维数组
x_np_2d = np.column_stack((pt_np, phi_np))
# 转Torch张量
x_torch = torch.from_numpy(x_np_2d)

标签y的转换 #

标签为一维数组，直接转换即可：

y_torch = torch.tensor(ak.to_numpy(y))

内容的提问来源于stack exchange，提问作者Teddy Curtis