PyTorch torch.unbind 函数
Pytorch torch 参考手册torch.unbind 是 PyTorch 中用于沿指定维度分割张量为元组的函数。
函数定义
torch.unbind(input, dim=0)
使用示例
实例
import torch
x = torch.arange(12).reshape(3, 4)
print("原始张量:")
print(x)
result = torch.unbind(x, dim=0)
print("沿 dim=0 分割:")
for i, t in enumerate(result):
print(f" 块 {i}: {t}")
result = torch.unbind(x, dim=1)
print("沿 dim=1 分割:")
for i, t in enumerate(result):
print(f" 块 {i}: {t}")
x = torch.arange(12).reshape(3, 4)
print("原始张量:")
print(x)
result = torch.unbind(x, dim=0)
print("沿 dim=0 分割:")
for i, t in enumerate(result):
print(f" 块 {i}: {t}")
result = torch.unbind(x, dim=1)
print("沿 dim=1 分割:")
for i, t in enumerate(result):
print(f" 块 {i}: {t}")
输出结果为：
原始张量:
tensor([[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]])
沿 dim=0 分割:
块 0: tensor([0, 1, 2, 3])
块 1: tensor([4, 5, 6, 7])
块 2: tensor([ 8, 9, 10, 11])
沿 dim=1 分割:
块 0: tensor([0, 4, 8])
块 1: tensor([1, 5, 9])
块 2: tensor([ 2, 6, 10])
块 3: tensor([ 3, 7, 11])