PyTorch 中如何使用 split 方法？

推荐答案

在 PyTorch 中，split 方法用于将张量沿着指定的维度分割成多个子张量。split 方法有两种使用方式：

1. 按大小分割：指定每个子张量的大小。
2. 按数量分割：指定分割的子张量数量，每个子张量大小相等。

按大小分割

// www.javascriptcn.com code example
import torch

# 创建一个张量
x = torch.arange(10).reshape(2, 5)
print(x)

# 按大小分割
result = torch.split(x, 2, dim=1)
print(result)

按数量分割

// www.javascriptcn.com code example
import torch

# 创建一个张量
x = torch.arange(10).reshape(2, 5)
print(x)

# 按数量分割
result = torch.split(x, 2, dim=1)
print(result)

本题详细解读

split 方法的基本用法

split 方法的语法如下：

torch.split(tensor, split_size_or_sections, dim=0)

• tensor: 要分割的张量。
• split_size_or_sections: 可以是一个整数或一个列表。如果是整数，表示每个子张量的大小；如果是列表，表示每个子张量的大小。
• dim: 指定沿着哪个维度进行分割，默认为 0。

按大小分割

当 split_size_or_sections 是一个整数时，split 方法会按照指定的大小沿着指定维度分割张量。例如：

x = torch.arange(10).reshape(2, 5)
result = torch.split(x, 2, dim=1)

在这个例子中，x 是一个形状为 (2, 5) 的张量，split 方法沿着第 1 维度（列）将其分割成大小为 2 的子张量。结果是一个包含 3 个子张量的元组，其中前两个子张量的大小为 (2, 2)，最后一个子张量的大小为 (2, 1)。

按数量分割

当 split_size_or_sections 是一个列表时，split 方法会按照列表中的大小沿着指定维度分割张量。例如：

x = torch.arange(10).reshape(2, 5)
result = torch.split(x, [2, 3], dim=1)

在这个例子中，x 是一个形状为 (2, 5) 的张量，split 方法沿着第 1 维度（列）将其分割成大小为 (2, 2) 和 (2, 3) 的两个子张量。

注意事项

• 如果 split_size_or_sections 是一个整数，且指定的分割大小不能整除张量的对应维度大小，最后一个子张量的大小会小于指定的分割大小。
• 如果 split_size_or_sections 是一个列表，列表中的元素之和必须等于张量对应维度的大小。

通过 split 方法，可以方便地将张量分割成多个子张量，这在处理批量数据或进行并行计算时非常有用。