在使用%ign%ignore_a_1%re_a_1%的`torch.vmap`进行函数向量化时，如果在函数内部创建新的张量（如通过`torch.zeros`），并且该张量的形状不完全由批处理输入直接决定，可能会遇到`batchedtensor`兼容性问题。本文将深入探讨这一挑战，并提供一种优雅的解决方案，通过结合`clone()`和`torch.concatenate`来确保内部创建的张量也能正确地进行批处理，从而实现高效且可读的代码。
理解 torch.vmap 与内部张量创建的挑战
torch.vmap 是 PyTorch 提供的一个强大工具，用于自动向量化函数，将处理单个样本的函数转换为处理一批样本的函数，而无需手动添加批处理维度。它的核心机制是追踪对 BatchedTensor（即 vmap 传入的批处理输入）进行的操作，并确保所有派生张量都正确地保留批处理维度。
然而，当我们在 vmap 内部的函数中，使用 torch.zeros 或 torch.empty 等函数创建一个新的张量时，如果这个张量的形状不是直接从 BatchedTensor 派生而来，或者我们尝试对其进行原地修改，vmap 可能无法正确识别并对其进行批处理。
考虑以下计算多项式伴随矩阵的函数 polycompanion：
import torchdef polycompanion(polynomial): # polynomial 是一个一维张量，表示多项式系数 # 例如 [a_n, a_{n-1}, ..., a_0] # 伴随矩阵的维度通常是 (degree, degree) # 这里的 deg = polynomial.shape[-1] - 2 是根据具体多项式定义来确定伴随矩阵维度的 deg = polynomial.shape[-1] - 2 # 初始化一个零矩阵作为伴随矩阵的基础 companion = torch.zeros((deg + 1, deg + 1)) # 填充伴随矩阵的左下角部分为单位矩阵 companion[1:, :-1] = torch.eye(deg) # 填充伴随矩阵的最后一列 # 这一列的值依赖于输入的 polynomial companion[:, -1] = -1. * polynomial[:-1] / polynomial[-1] return companion# 准备批处理输入poly_batched = torch.tensor([[1, 2, 3, 4], [1, 2, 3, 4]], dtype=torch.float32)# 尝试使用 vmap 向量化函数polycompanion_vmap = torch.vmap(polycompanion)# 运行会失败，因为 companion 在函数内部被创建为普通的 Tensor，# 并且在 vmap 运行时没有被正确地批处理。try: print(polycompanion_vmap(poly_batched))except Exception as e: print(f"Error when trying to vmap with internal torch.zeros: {e}")
上述代码在 torch.vmap 环境下运行时会抛出错误。原因是 companion = torch.zeros((deg+1, deg+1)) 创建了一个普通的 torch.Tensor，而不是一个 BatchedTensor。尽管 deg 是从 polynomial 的形状推导而来，但 torch.zeros 本身并不知道它需要被批处理。随后的原地赋值操作 companion[1:,:-1] = torch.eye(deg) 和 companion[:,-1] = … 并没有将 companion 转换为 BatchedTensor，导致 vmap 无法正确处理其批处理维度。
一个常见的临时解决方案是，在调用 vmap 之前预先创建并传入一个批处理的零张量：
# 初始的丑陋 workarounddef polycompanion_workaround(polynomial, companion_template): # deg = companion_template.shape[-1] - 1 # 这里的deg需要根据template重新计算 deg = polynomial.shape[-1] - 2 # 保持与原函数一致 # 使用传入的模板进行原地修改 companion_template[1:,:-1] = torch.eye(deg) companion_template[:,-1] = -1. * polynomial[:-1] / polynomial[-1] return companion_templatepolycompanion_vmap_workaround = torch.vmap(polycompanion_workaround)# 预先创建批处理的零张量# 注意这里的形状需要与期望的批处理输出形状匹配batch_size = poly_batched.shape[0]matrix_dim = poly_batched.shape[-1] - 1 # (deg+1)pre_batched_zeros = torch.zeros(batch_size, matrix_dim, matrix_dim, dtype=torch.float32)print("Workaround output:")print(polycompanion_vmap_workaround(poly_batched, pre_batched_zeros))
这种方法虽然能工作，但它要求修改函数签名以接受一个额外的模板张量，并且需要在外部手动管理这个模板的批处理维度，这降低了代码的优雅性和可维护性。
ProWritingAid
AI写作助手软件
114 查看详情
推荐解决方案：结合 clone() 和 torch.concatenate
为了在 vmap 内部优雅地创建并填充批处理张量，我们应该避免对非批处理张量进行原地修改。相反，我们可以通过组合已经批处理的张量片段来构建最终的批处理张量。这里的关键在于利用 clone() 操作来“激活”批处理维度，并使用 torch.concatenate 将这些片段组合起来。
以下是改进后的 polycompanion 函数：
import torchdef polycompanion_optimized(polynomial): deg = polynomial.shape[-1] - 2 # 1. 创建一个非批处理的零矩阵作为基础结构 # 这个 companion 此时仍是普通的 Tensor companion_base = torch.zeros((deg + 1, deg + 1), dtype=polynomial.dtype) # 2. 填充单位矩阵部分 # 这一部分是伴随矩阵的左侧部分 identity_part = companion_base[1:, :-1].clone() # 关键：clone() 使得这部分成为 BatchedTensor identity_part[torch.eye(deg, dtype=torch.bool)] = 1.0 # 填充单位矩阵 # 3. 计算伴随矩阵的最后一列 # polynomial 是 BatchedTensor，所以这个计算结果自然也是 BatchedTensor last_column = -1. * polynomial[:-1] / polynomial[-1] # 4. 扩展 last_column 的维度以匹配 concatenate 的要求 # last_column 的形状是 (deg+1,)，需要变成 (deg+1, 1) 才能与 identity_part 拼接 last_column_expanded = last_column[:, None] # 5. 使用 torch.concatenate 将批处理的片段组合起来 # identity_part 是 (deg+1, deg) 形状的 BatchedTensor # last_column_expanded 是 (deg+1, 1) 形状的 BatchedTensor # 沿着 dim=1 拼接，得到 (deg+1, deg+1) 形状的 BatchedTensor _companion = torch.concatenate([identity_part, last_column_expanded], dim=1) return _companion# 准备批处理输入poly_batched = torch.tensor([[1, 2, 3, 4], [1, 2, 3, 4]], dtype=torch.float32)# 使用 vmap 向量化优化后的函数polycompanion_vmap_optimized = torch.vmap(polycompanion_optimized)print("nOptimized vmap output:")print(polycompanion_vmap_optimized(poly_batched))
输出结果：
Optimized vmap output:tensor([[[ 0.0000, 0.0000, -0.2500], [ 1.0000, 0.0000, -0.5000], [ 0.0000, 1.0000, -0.7500]], [[ 0.0000, 0.0000, -0.2500], [ 1.0000, 0.0000, -0.5000], [ 0.0000, 1.0000, -0.7500]]])
解决方案解析
companion_base = torch.zeros((deg + 1, deg + 1), dtype=polynomial.dtype): 我们仍然可以创建一个普通的零矩阵作为基础，用于确定形状。重要的是，我们不直接对它进行原地修改并期望它被批处理。identity_part = companion_base[1:, :-1].clone(): 这是关键一步。我们首先从 companion_base 中切片出需要填充单位矩阵的部分。然后，通过调用 .clone()，即使 companion_base 本身不是 BatchedTensor，这个 identity_part 在 vmap 内部执行时，会被 vmap 机制识别为一个需要批处理的张量。它会为每个批次创建一个独立的副本。identity_part[torch.eye(deg, dtype=torch.bool)] = 1.0: 在 clone() 之后，我们可以在这个 identity_part 上进行原地修改（填充单位矩阵），因为 vmap 已经将其视为一个批处理的张量。*`last_column = -1. polynomial[:-1] / polynomial[-1]**: 这一部分直接从polynomial（一个BatchedTensor）计算而来，因此last_column自然也是一个BatchedTensor`。last_column_expanded = last_column[:, None]: 使用 None 或 unsqueeze(-1) 扩展 last_column 的维度，使其形状从 (deg+1,) 变为 (deg+1, 1)，以便与 identity_part 沿着 dim=1 进行拼接。_companion = torch.concatenate([identity_part, last_column_expanded], dim=1): 最后，我们将两个已经正确批处理的张量片段 identity_part 和 last_column_expanded 沿着正确的维度拼接起来，形成最终的批处理伴随矩阵 _companion。
注意事项与最佳实践
理解 vmap 的批处理传播机制：vmap 通过跟踪对 BatchedTensor 的操作来自动传播批处理维度。如果一个张量在函数内部被创建，但其创建过程没有直接依赖于 BatchedTensor 或没有通过 clone() 等方式“激活”批处理，它将不会被 vmap 自动批处理。避免在 vmap 内部对非批处理张量进行原地修改：尽量避免直接创建 torch.zeros 或 torch.empty 然后对其进行原地修改，除非你确定这些操作不会影响 vmap 的批处理行为，或者你正在修改一个已经被 vmap 识别为 BatchedTensor 的张量。利用 clone() 激活批处理：当需要从一个非批处理的张量结构中派生出需要批处理的部分时，使用 .clone() 是一个有效的策略。它会为每个批次创建一个独立的副本，并允许 vmap 追踪这些副本。构建而非修改：在 vmap 内部，更推荐通过组合已批处理的张量片段来构建目标张量，而不是尝试原地修改一个最初未批处理的张量。torch.zeros_like 或 torch.empty_like：如果新张量的形状与某个 BatchedTensor 输入完全一致，使用 torch.zeros_like(input_tensor) 或 torch.empty_like(input_tensor) 可以直接创建批处理的零/空张量，这是最简单直接的方法。但在本例中，伴随矩阵的形状与 polynomial 的形状不同，所以不能直接使用 zeros_like。
总结
在 torch.vmap 中进行动态张量创建时，关键在于确保所有需要批处理的中间张量都能够被 vmap 机制正确识别。通过将张量的创建分解为批处理的片段（例如，利用 clone() 激活批处理维度，或直接从 BatchedTensor 派生），然后使用 torch.concatenate 等操作将它们组合起来，我们可以编写出既高效又符合 vmap 设计哲学的代码。这种方法避免了修改函数签名来传入预批处理张量的“丑陋”做法，使得代码更加清晰和易于维护。
以上就是PyTorch vmap中动态张量创建的技巧与最佳实践的详细内容，更多请关注创想鸟其它相关文章！
版权声明：本文内容由互联网用户自发贡献，该文观点仅代表作者本人。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。
如发现本站有涉嫌抄袭侵权/违法违规的内容， 请发送邮件至 chuangxiangniao@163.com 举报，一经查实，本站将立刻删除。
发布者：程序猿，转转请注明出处：https://www.chuangxiangniao.com/p/918574.html
微信扫一扫 
支付宝扫一扫 
