PyTorch torch.vsplit Function
Pytorch torch Reference Manual
torch.vsplitis a function in PyTorch used to split tensors vertically (along rows).
Function Definition
torch.vsplit(input, indices_or_sections)
Usage Example
Example
import torch
# 2D tensor vertical split
x = torch.arange(12).reshape(4, 3)
print("Original 2D tensor:")
print(x)
result = torch.vsplit(x, 2)
print("Split evenly into 2 parts along rows:")
for i, t in enumerate(result):
print(f" Block {i}:n{t}")
# Split by index
result = torch.vsplit(x, [1, 3])
print("nSplit by index [1, 3]:")
for i, t in enumerate(result):
print(f" Block {i}:n{t}")
# 3D tensor vertical split
y = torch.arange(24).reshape(4, 3, 2)
print("n3D tensor:")
print(y)
result = torch.vsplit(y, 2)
print("nSplit into 2 parts along the first dimension:")
for i, t in enumerate(result):
print(f" Block {i}: {t.shape}")
# 2D tensor vertical split
x = torch.arange(12).reshape(4, 3)
print("Original 2D tensor:")
print(x)
result = torch.vsplit(x, 2)
print("Split evenly into 2 parts along rows:")
for i, t in enumerate(result):
print(f" Block {i}:n{t}")
# Split by index
result = torch.vsplit(x, [1, 3])
print("nSplit by index [1, 3]:")
for i, t in enumerate(result):
print(f" Block {i}:n{t}")
# 3D tensor vertical split
y = torch.arange(24).reshape(4, 3, 2)
print("n3D tensor:")
print(y)
result = torch.vsplit(y, 2)
print("nSplit into 2 parts along the first dimension:")
for i, t in enumerate(result):
print(f" Block {i}: {t.shape}")
The output result is:
原始二维张量:
tensor([[ 0, 1, 2],
[ 3, 4, 5],
[ 6, 7, 8],
[ 9, 10, 11]])
沿行平均分为 2 份:
块 0:
tensor([[0, 1, 2],
[3, 4, 5]])
Chunk 1:
tensor([[ 6, 7, 8],
[ 9, 10, 11]])
按索引 [1, 3] 分割:
块 0:
tensor([[0, 1, 2]])
Chunk 1:
tensor([[3, 4, 5],
[6, 7, 8]])
Chunk 2:
tensor([[ 9, 10, 11]])
三维张量:
tensor([[[ 0, 1],
[ 3, 4],
[ 5, 6]],
[[ 7, 8],
[ 9, 10],
[11, 12]],
[[13, 14],
[15, 16],
[17, 18]],
[[19, 20],
[21, 22],
[23, 24]]])
沿第一维分为 2 份:
块 0: torch.Size([2, 3, 2])
Chunk 1: torch.Size([2, 3, 2])
Other Extensions