案例:
import torch
a = torch.rand(size=[3,4])
print(a)
out = torch.chunk(a, 2, dim=0)
print(out)
a = torch.rand(size=[3,4])
print(a)
out = torch.split(a, 2, dim=0)
print(out)
a = torch.rand(size=[6,4])
print(a)
out = torch.split(a, [2,4], dim=0)
print(out)
"""
out:
tensor([[0.4318, 0.9049, 0.8166, 0.3588],
[0.6499, 0.5793, 0.9395, 0.6760],
[0.2342, 0.9079, 0.2751, 0.6862]])
(tensor([[0.4318, 0.9049, 0.8166, 0.3588],
[0.6499, 0.5793, 0.9395, 0.6760]]), tensor([[0.2342, 0.9079, 0.2751, 0.6862]]))
tensor([[0.2170, 0.0928, 0.9245, 0.3909],
[0.6361, 0.4667, 0.5978, 0.4724],
[0.1860, 0.4510, 0.1610, 0.9921]])
(tensor([[0.2170, 0.0928, 0.9245, 0.3909],
[0.6361, 0.4667, 0.5978, 0.4724]]), tensor([[0.1860, 0.4510, 0.1610, 0.9921]]))
tensor([[0.6844, 0.9560, 0.7572, 0.0400],
[0.4103, 0.3539, 0.1543, 0.5854],
[0.7305, 0.9435, 0.8111, 0.6618],
[0.5152, 0.9562, 0.2030, 0.3958],
[0.3747, 0.8432, 0.9954, 0.4090],
[0.6239, 0.3010, 0.6631, 0.2251]])
(tensor([[0.6844, 0.9560, 0.7572, 0.0400],
[0.4103, 0.3539, 0.1543, 0.5854]]), tensor([[0.7305, 0.9435, 0.8111, 0.6618],
[0.5152, 0.9562, 0.2030, 0.3958],
[0.3747, 0.8432, 0.9954, 0.4090],
[0.6239, 0.3010, 0.6631, 0.2251]]))
Process finished with exit code 0
"""