当前位置: 首页>>代码示例 >>用法及示例精选 >>正文


Python PyTorch tensor_split用法及代码示例


本文简要介绍python语言中 torch.tensor_split 的用法。

用法:

torch.tensor_split(input, indices_or_sections, dim=0) → List of Tensors

参数

  • input(Tensor) -要分裂的张量

  • indices_or_sections(Tensor,int或者list或者python的元组:ints) -

    如果 indices_or_sections 是整数 n 或值为 n 的零维长张量,则 input 沿维度 dim 拆分为 n 部分。如果 input 沿维度 dim 可被 n 整除,则每个部分将具有相同的大小 input.size(dim) / n 。如果 input 不能被 n 整除,则第一个 int(input.size(dim) % n) 部分的大小将为 int(input.size(dim) / n) + 1 ,其余部分的大小为 int(input.size(dim) / n)

    如果indices_or_sections 是整数列表或元组,或一维长张量,则input 在列表、元组或张量中的每个索引处沿维度dim 拆分。例如,indices_or_sections=[2, 3]dim=0 将产生张量 input[:2]input[2:3]input[3:]

    如果indices_or_sections是张量,那么它在CPU上必须是零维或一维长张量。

  • dim(int,可选的) -沿其分割张量的维度。默认值:0

根据 indices_or_sections 指定的索引或部分数量,将张量拆分为多个 sub-tensors,所有这些都是 input 的视图,沿维度 dim。该函数基于 NumPy 的 numpy.array_split()

例子:

>>> x = torch.arange(8)
>>> torch.tensor_split(x, 3)
(tensor([0, 1, 2]), tensor([3, 4, 5]), tensor([6, 7]))

>>> x = torch.arange(7)
>>> torch.tensor_split(x, 3)
(tensor([0, 1, 2]), tensor([3, 4]), tensor([5, 6]))
>>> torch.tensor_split(x, (1, 6))
(tensor([0]), tensor([1, 2, 3, 4, 5]), tensor([6]))

>>> x = torch.arange(14).reshape(2, 7)
>>> x
tensor([[ 0,  1,  2,  3,  4,  5,  6],
        [ 7,  8,  9, 10, 11, 12, 13]])
>>> torch.tensor_split(x, 3, dim=1)
(tensor([[0, 1, 2],
        [7, 8, 9]]),
 tensor([[ 3,  4],
        [10, 11]]),
 tensor([[ 5,  6],
        [12, 13]]))
>>> torch.tensor_split(x, (1, 6), dim=1)
(tensor([[0],
        [7]]),
 tensor([[ 1,  2,  3,  4,  5],
        [ 8,  9, 10, 11, 12]]),
 tensor([[ 6],
        [13]]))

相关用法


注:本文由纯净天空筛选整理自pytorch.org大神的英文原创作品 torch.tensor_split。非经特殊声明,原始代码版权归原作者所有,本译文未经允许或授权,请勿转载或复制。