pytorch中?torch.unbind、torch.split 和 torch.chunk等函數(shù)可用于張量的拆分操作。

1. torch.unbind

功能說明：
torch.unbind?沿指定的維度將張量“解包”為多個張量，返回一個元組。解包后被操作的那個維度會消失，每個輸出張量的維度數(shù)會比原張量少 1。

函數(shù)簽名：

torch.unbind(input, dim=0)

• input: 待解包的張量。
• dim: 指定沿哪個維度解包。

多維張量示例：

import torch# 創(chuàng)建一個形狀為 (2, 3, 4) 的張量
x = torch.arange(24).reshape(2, 3, 4)
print("原始張量 x:\n", x)
print("x.shape:", x.shape)  # torch.Size([2, 3, 4])# 沿第 0 維解包
slices0 = torch.unbind(x, dim=0)
print("\n沿 dim=0 解包：")
for i, t in enumerate(slices0):print(f"slice {i} shape: {t.shape}")# 每個張量形狀為 (3, 4)# 沿第 1 維解包
slices1 = torch.unbind(x, dim=1)
print("\n沿 dim=1 解包：")
for i, t in enumerate(slices1):print(f"slice {i} shape: {t.shape}")# 每個張量形狀為 (2, 4)

2. torch.split

功能說明：
torch.split?根據(jù)給定的大小或尺寸列表，將張量沿指定維度切分成若干塊。

• 如果傳入一個整數(shù)，則每塊的大小為該整數(shù)，最后一塊可能會小于這個整數(shù)。
• 如果傳入一個尺寸列表，則按列表中指定的尺寸進行切分。

函數(shù)簽名：

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

• tensor: 待分割的張量。
• split_size_or_sections: 整數(shù)或尺寸列表，指定每塊的大小。
• dim: 指定沿哪個維度進行切分。

多維張量示例：

import torch# 創(chuàng)建一個形狀為 (2, 5, 4) 的張量
x = torch.arange(40).reshape(2, 5, 4)
print("原始張量 x:\n", x)
print("x.shape:", x.shape)  # torch.Size([2, 5, 4])# 按照固定大小進行切分：沿第 1 維，每塊大小為 2
splits_fixed = torch.split(x, 2, dim=1)
print("\n沿 dim=1 按固定大小 2 切分：")
for i, t in enumerate(splits_fixed):print(f"chunk {i} shape: {t.shape}")# 輸出塊的形狀可能為 (2, 2, 4), (2, 2, 4) 和最后一塊 (2, 1, 4)# 按照指定尺寸列表進行切分：沿第 1 維，分塊尺寸為 [1, 2, 2]
splits_list = torch.split(x, [1, 2, 2], dim=1)
print("\n沿 dim=1 按尺寸列表 [1, 2, 2] 切分：")
for i, t in enumerate(splits_list):print(f"chunk {i} shape: {t.shape}")# 分別輸出形狀 (2, 1, 4), (2, 2, 4), (2, 2, 4)

3. torch.chunk

功能說明：
torch.chunk?將張量沿指定維度平均分成指定數(shù)量的塊。如果張量在該維度上的長度不能被塊數(shù)整除，則前面的塊會比后面塊多一個元素（塊的尺寸差別最多為 1）。

函數(shù)簽名：

torch.chunk(tensor, chunks, dim=0)

• tensor: 待分割的張量。
• chunks: 指定分成幾塊。
• dim: 指定沿哪個維度進行分塊。

多維張量示例：

對比總結

函數(shù)	分割方式	返回結果形式	適用場景
torch.unbind	沿指定維度將張量完全解包，每個輸出不含該維度	元組，輸出張量數(shù) = 該維度的長度	需要逐個處理某一維度上的切片，且希望移除該維度時使用。
torch.split	按照指定大小或尺寸列表切分張量	元組或列表	需要按固定大小或自定義尺寸列表切分張量，最后一塊可能不均勻。
torch.chunk	將張量均勻分成指定數(shù)量的塊	元組或列表	希望將張量平均分成若干塊，塊數(shù)固定，自動處理無法整除的情況。

注意：

• 當處理多維張量時，選擇沿哪一維進行分割非常重要；
• torch.unbind?會移除分割的那個維度，而?torch.split?和?torch.chunk?則保持原始維度，只是該維度上的大小發(fā)生變化。

通過這些示例代碼和說明，你可以根據(jù)具體需求選擇合適的函數(shù)來分割多維張量。