megengine.functional.flatten

flatten(inp, start_axis=0, end_axis=- 1)[源代码]

通过将子张量从 start_axis 维展平到 end_axis 维,实现对张量的重塑(reshape)。

参数
  • inp – 输入张量。

  • start_axis – 子张量被展平时的初始维数。 默认: 0

  • end_axis – 子张量被展平时的最终维数。 默认: -1

返回

输出张量。

例如:

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

inp_shape = (2, 2, 3, 3)
x = tensor(
    np.arange(36, dtype=np.int32).reshape(inp_shape),
)
out = F.flatten(x, 2)
print(x.numpy().shape)
print(out.numpy().shape)

输出:

(2, 2, 3, 3)
(2, 2, 9)