megengine.functional.split

split(inp, nsplits_or_sections, axis=0)[源代码]

把一个张量分隔成很多个小张量。当`nsplits_or_sections`是整数时,最后一个张量可能比其他张量小。

参数
  • inp – 输入张量。

  • nsplits_or_sections – 子张量或每个部分的信息列表的个数。

  • axis – 需要被分隔的轴。

返回

输出张量列表。

例如:

import os
import numpy as np
from megengine import tensor
import megengine.functional as F

x = tensor(np.random.random((10, 20)), dtype=np.float32)
y = F.split(x, 3)
z = F.split(x, [6, 17], axis=1)

print([i.numpy().shape for i in y])
print([i.numpy().shape for i in z])

输出:

[(4, 20), (3, 20), (3, 20)]
[(10, 6), (10, 11), (10, 3)]