（一）PyTorch 张量

Tensor 数据类型

Data type	dtype	CPU Tensor	GPU Tensor
Boolean	torch.bool	torch.BoolTensor	torch.cuda.BoolTensor
8-bit integer (unsigned)	torch.uint8	torch.ByteTensor	torch.cuda.ByteTensor
8-bit integer (signed)	torch.int8	torch.CharTensor	torch.cuda.CharTensor
16-bit integer (signed)	torch.int16 or torch.short	torch.ShortTensor	torch.cuda.ShortTensor
32-bit integer (signed)	torch.int32 or torch.int	torch.IntTensor	torch.cuda.IntTensor
64-bit integer (signed)	torch.int64 or torch.long	torch.LongTensor	torch.cuda.LongTensor
16-bit floating point	torch.float16 or torch.half	torch.HalfTensor	torch.cuda.HalfTensor
16-bit floating point	torch.bfloat16	torch.BFloat16Tensor	torch.cuda.BFloat16Tensor
32-bit floating point	torch.float32 or torch.float	torch.FloatTensor	torch.cuda.FloatTensor
64-bit floating point	torch.float64 or torch.double	torch.DoubleTensor	torch.cuda.DoubleTensor

设置 Tensor 默认类型

import torch

print(torch.tensor(True).dtype)  # torch.bool
print(torch.tensor(True).type())  # torch.BoolTensor
print(torch.tensor(1).dtype)  # torch.int64
print(torch.tensor(1).type())  # torch.LongTensor
print(torch.tensor(1.).dtype)  # torch.float32
print(torch.tensor(1.).type())  # torch.FloatTensor
print(torch.tensor(3j).dtype)  # torch.complex64
print(torch.tensor(3j).type())  # torch.ComplexFloatTensor

torch.set_default_dtype(torch.double)  # 设置默认类型为 double

print(torch.tensor(True).dtype)  # torch.bool
print(torch.tensor(True).type())  # torch.BoolTensor
print(torch.tensor(1).dtype)  # torch.int64
print(torch.tensor(1).type())  # torch.LongTensor
print(torch.tensor(1.).dtype)  # torch.float64
print(torch.tensor(1.).type())  # torch.DoubleTensor
print(torch.tensor(3j).dtype)  # torch.complex128
print(torch.tensor(3j).type())  # torch.ComplexDoubleTensor

set_default_dtype()只能设置floating-point类型，否则会报TypeError: only floating-point types are supported as the default type错误。

标量与张量

import torch

# 标量
a = torch.tensor(1)
print(a.shape)  # torch.Size([])
print(a.size())  # torch.Size([])
print(a.ndim)  # 0
print(a.dim())  # 0

# 一维张量
b = torch.tensor([1, 2])
print(b.shape)  # torch.Size([2])
print(b.size())  # torch.Size([2])
print(b.ndim)  # 1
print(b.dim())  # 1

# 二维张量
c = torch.tensor([[1, 2]])
print(c.shape)  # torch.Size([1, 2])
print(c.size())  # torch.Size([1, 2])
print(c.ndim)  # 2
print(c.dim())  # 2

标量是一个单独的数，ndim为0。

创建 Tensor

.tensor

import torch

# 标量
print(torch.tensor(1))  # tensor(1)
print(torch.tensor(1, dtype=torch.float64))  # tensor(1., dtype=torch.float64)

# 张量
print(torch.tensor([1, 2]))  # tensor([1, 2])
print(torch.tensor([1, 2], dtype=torch.float64))  # tensor([1., 2.], dtype=torch.float64)

.from_numpy

import numpy as np
import torch

data = np.array([1, 2, 3])
print(torch.from_numpy(data))  # tensor([1, 2, 3], dtype=torch.int32)
data = np.array([1, 2, 3], dtype=np.float64)
print(torch.from_numpy(data))  # tensor([1., 2., 3.], dtype=torch.float64)

Tensor

import torch

# 参数为 shape 大小
print(torch.Tensor(1))  # tensor([-3.0434e+31])
print(torch.Tensor(1, 2))  # tensor([[0., 0.]])
print(torch.Tensor(1, 2, 3))
"""
tensor([[[-3.0741e+31,  1.6031e-42,  0.0000e+00],
         [ 0.0000e+00,  0.0000e+00,  0.0000e+00]]])
"""

# 参数为列表
print(torch.Tensor([1]))  # tensor([1.])
print(torch.Tensor([1, 2]))  # tensor([1., 2.])
print(torch.Tensor([1, 2, 3]))  # tensor([1., 2., 3.])

Tensor支持两种传参方式：

1. 当参数为列表时，创建列表对应维度的Tensor并初始化数据为列表数据。
2. 当参数不为列表时，与.empty()类似，创建参数指定的shape的空的Tensor。

BoolTensor、ByteTensor、CharTensor、ShortTensor、IntTensor、LongTensor、HalfTensor、FloatTensor、DoubleTensor也是一样。

.empty/.zeros/.ones/.full/.eye

import torch

input = torch.tensor([1, 1, 1, 1, 1])

print(torch.empty(()))  # tensor(-6.5391e-19)
print(torch.empty((1, 5)))  # tensor([[-6.6069e-19,  1.3943e-42,  0.0000e+00,  0.0000e+00,  0.0000e+00]])
print(torch.empty_like((input)))  # tensor([0, 0, 0, 0, 0])

print(torch.zeros(()))  # tensor(0.)
print(torch.zeros((1, 5)))  # tensor([[0., 0., 0., 0., 0.]])
print(torch.zeros_like((input)))  # tensor([0, 0, 0, 0, 0])

print(torch.ones(()))  # tensor(1.)
print(torch.ones((1, 5)))  # tensor([[1., 1., 1., 1., 1.]])
print(torch.ones_like((input)))  # tensor([1, 1, 1, 1, 1])

print(torch.full((), 100))  # tensor(100)
print(torch.full((1, 5), 100))  # tensor([[100, 100, 100, 100, 100]])
print(torch.full_like((input), 100))  # tensor([100, 100, 100, 100, 100])

print(torch.eye(3))
"""
tensor([[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]])
"""
print(torch.eye(3, 5))
"""
tensor([[1., 0., 0., 0., 0.],
        [0., 1., 0., 0., 0.],
        [0., 0., 1., 0., 0.]])
"""

.arange/.linspace/.logspace

import torch

print(torch.arange(0, 10))  # tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
print(torch.arange(0, 10, step=3))  # tensor([0, 3, 6, 9])

# [0, 1] 等分成 5 份
print(torch.linspace(0, 1, steps=5))  # tensor([0.0000, 0.2500, 0.5000, 0.7500, 1.0000])

# 默认以 10 为底数，[0, 1] 等分成 5 份做为指数
print(torch.logspace(0, 1, steps=5))  # tensor([ 1.0000,  1.7783,  3.1623,  5.6234, 10.0000])
# 以 2 为底数，[0, 1] 等分成 5 份做为指数
print(torch.logspace(0, 1, steps=5, base=2))  # tensor([1.0000, 1.1892, 1.4142, 1.6818, 2.0000])

随机采样

随机种子

import torch

# 设置 CPU 随机种子
torch.manual_seed(1)

# 设置 GPU 随机种子
torch.cuda.manual_seed(1)

# 查看设置的随机种子
print(torch.initial_seed())  # 1

# 随机设置随机种子
torch.seed()

print(torch.initial_seed())  # 22287915889500

随机函数

.rand/.rand_like

import torch

# 固定随机种子
torch.manual_seed(1)

input = torch.empty(3, 3)

print(torch.rand(3, 3))
"""
tensor([[0.7576, 0.2793, 0.4031],
        [0.7347, 0.0293, 0.7999],
        [0.3971, 0.7544, 0.5695]])
"""

print(torch.rand_like(input))
"""
tensor([[0.4388, 0.6387, 0.5247],
        [0.6826, 0.3051, 0.4635],
        [0.4550, 0.5725, 0.4980]])
"""

torch.rand()返回在区间[0, 1)均匀分布的随机数填充的张量。

.randint/.randint_like

import torch

# 固定随机种子
torch.manual_seed(1)

input = torch.empty(3, 3)

print(torch.randint(low=0, high=10, size=(3, 3)))
"""
tensor([[5, 9, 4],
        [8, 3, 3],
        [1, 1, 9]])
"""

print(torch.randint_like(input, high=10))
"""
tensor([[2., 8., 9.],
        [6., 3., 3.],
        [0., 2., 1.]])
"""

torch.randint(low, high)返回在区间[low, high)的随机数填充的张量。

.randperm

import torch

# 固定随机种子
torch.manual_seed(1)

print(torch.randperm(10))  # tensor([5, 6, 1, 2, 0, 8, 9, 3, 7, 4])

torch.randperm(n)返回在区间[0, n)的随机排列整数。

.randn/.randn_like

import torch

# 固定随机种子
torch.manual_seed(1)

input = torch.empty(3, 3)

# 标准正态分布 N(0, 1)，均值为 0，方差为 1
print(torch.randn(3, 3))
"""
tensor([[ 0.6614,  0.2669,  0.0617],
        [ 0.6213, -0.4519, -0.1661],
        [-1.5228,  0.3817, -1.0276]])
"""

print(torch.randn_like(input))
"""
tensor([[-0.5631, -0.8923, -0.0583],
        [-0.1955, -0.9656,  0.4224],
        [ 0.2673, -0.4212, -0.5107]])
"""

torch.randn()从标准正态分布中随机采样。

.normal

import torch

# 固定随机种子
torch.manual_seed(1)

# 离散正态分布 N(mean, std)
print(torch.normal(mean=torch.full((5,), 0.), std=torch.arange(0, 1, 0.2)))  # tensor([ 0.0000,  0.0534,  0.0247,  0.3728, -0.3615])

torch.normal(mean, std)从给定参数mean、std的离散正态分布中随机采样。

.bernoulli

import torch

input = torch.empty(3, 3).uniform_(0, 1)  # generate a uniform random matrix with range [0, 1]

print(torch.bernoulli(input))
"""
tensor([[1., 0., 0.],
        [1., 0., 1.],
        [0., 1., 1.]])
"""

print(torch.bernoulli(torch.ones(3, 3)))  # probability of drawing "1" is 1
"""
tensor([[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]])
"""

print(torch.bernoulli(torch.zeros(3, 3)))  # probability of drawing "1" is 0
"""
tensor([[0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.]])
"""

torch.bernoulli()从伯努利分布中抽取二进制随机数（0 或 1）。输入的值必须在[0, 1]范围内。

.poisson

import torch

torch.manual_seed(1)

rates = torch.rand(3, 3) * 5  # rate parameter between 0 and 5
print(torch.poisson(rates))
"""
tensor([[5., 1., 1.],
        [3., 0., 2.],
        [0., 2., 0.]])
"""

torch.poisson()从泊松分布中随机采样。

.multinomial

import torch

weights = torch.tensor([[1, 3, 5, 7, 9], [2, 4, 6, 8, 10]], dtype=torch.float)

print(torch.multinomial(weights, 5))
"""
tensor([[2, 3, 4, 1, 0],
        [3, 2, 1, 4, 0]])
"""

# replacement 默认为 False，表示不允许重复抽查，所以采样次数不能大于抽查个数
# print(torch.multinomial(weights, 6))  # RuntimeError: cannot sample n_sample > prob_dist.size(-1) samples without replacement

# replacement=True 允许重复抽查
print(torch.multinomial(weights, 6, replacement=True))
"""
tensor([[3, 0, 3, 3, 4, 2],
        [0, 1, 4, 4, 2, 4]])
"""

torch.multinomial(input, num_samples, replacement)对input的每一行做从多项式分布中采样num_samples次，输出的张量是每一次取值时input张量对应行的下标。

索引与切片

Python 语法

import torch

# 4 张 3 通道，高 28，宽 28 的图片
images = torch.rand(4, 3, 28, 28)
print(images.shape)  # torch.Size([4, 3, 28, 28])

# 获取第一张图片
print(images[0].shape)  # torch.Size([3, 28, 28])
# 获取第一张图片的第一个通道数据
print(images[0, 0].shape)  # torch.Size([28, 28])
# 获取前两张图片
print(images[:2].shape)  # torch.Size([2, 3, 28, 28])
# 获取前两张图片前两个通道数据
print(images[:2, :2].shape)  # torch.Size([2, 2, 28, 28])
# 获取前两张图片最后一个通道数据
print(images[:2, -1].shape)  # torch.Size([2, 28, 28])
# 间隔获取图片数据
print(images[:, :, ::2, ::2].shape)  # torch.Size([4, 3, 14, 14])

# 获取所有图片
print(images[...].shape)  # torch.Size([4, 3, 28, 28])
# 获取第一张图片
print(images[0, ...].shape)  # torch.Size([3, 28, 28])
# 间隔获取图片宽数据
print(images[..., ::2].shape)  # torch.Size([4, 3, 28, 14])

.narrow/.narrow_copy

import torch

data = torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9 ,10], [11, 12, 13, 14, 15]])
print(data)
"""
tensor([[ 1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10],
        [11, 12, 13, 14, 15]])
"""

print(torch.narrow(data, dim=0, start=1, length=2))
"""
tensor([[ 6,  7,  8,  9, 10],
        [11, 12, 13, 14, 15]])
"""
print(torch.narrow(data, dim=1, start=1, length=3))
"""
tensor([[ 2,  3,  4],
        [ 7,  8,  9],
        [12, 13, 14]])
"""

print(torch.narrow_copy(data, dim=0, start=1, length=2))
"""
tensor([[ 6,  7,  8,  9, 10],
        [11, 12, 13, 14, 15]])
"""
print(torch.narrow_copy(data, dim=1, start=1, length=3))
"""
tensor([[ 2,  3,  4],
        [ 7,  8,  9],
        [12, 13, 14]])
"""

torch.narrow()在指定维度缩小张量，可以简单理解为类似切片，tensor[start: start + length]

torch.narrow_copy()与torch.narrow()相同，但返回的是副本而不是共享存储。

.select/.index_select

import torch

# 4 张 3 通道，高 28，宽 28 的图片
images = torch.rand(4, 3, 28, 28)
print(images.shape)  # torch.Size([4, 3, 28, 28])

# 获取第一张图片
print(torch.select(images, dim=0, index=0).shape)  # torch.Size([3, 28, 28])
# 获取所有图片的第一个通道
print(torch.select(images, dim=1, index=1).shape)  # torch.Size([4, 28, 28])

# 获取第二张和第四张图片
print(torch.index_select(images, dim=0, index=torch.tensor([1, 3])).shape)  # torch.Size([2, 3, 28, 28])
# 获取所有图片二三通道数据
print(torch.index_select(images, dim=1, index=torch.tensor([1, 2])).shape)  # torch.Size([4, 2, 28, 28])
# 间隔获取所有图片高数据
print(torch.index_select(images, dim=2, index=torch.arange(0, 28, 2)).shape)  # torch.Size([4, 3, 14, 28])
# 间隔获取所有图片宽数据
print(torch.index_select(images, dim=3, index=torch.arange(0, 28, 2)).shape)  # torch.Size([4, 3, 28, 14])

torch.select()沿着维度dim，在给定索引index对input张量进行切片，等价于切片。比如：
tensor.select(0, index)等于tensor[index]。
tensor.select(2, index)等于tensor[:,:,index]。

torch.index_select()沿着维度dim对input张量进行索引。

.masked_select

import torch

torch.manual_seed(1)

data = torch.randn(3, 4)
print(data)
"""
tensor([[ 0.6614,  0.2669,  0.0617,  0.6213],
        [-0.4519, -0.1661, -1.5228,  0.3817],
        [-1.0276, -0.5631, -0.8923, -0.0583]])
"""

mask = data.ge(0)
print(mask)
"""
tensor([[ True,  True,  True,  True],
        [False, False, False,  True],
        [False, False, False, False]])
"""

print(torch.masked_select(data, mask))  # tensor([0.6614, 0.2669, 0.0617, 0.6213, 0.3817])

torch.masked_select()根据布尔掩码选择数据，返回的是一维数据。

.gather

import torch

torch.manual_seed(1)

data = torch.randperm(16).view(4,4)
print(data)
"""
tensor([[ 5, 15,  6,  4],
        [11,  2,  7, 12],
        [ 1,  0,  9,  8],
        [10,  3, 13, 14]])
"""

index = torch.tensor([[1, 0, 3, 2]])
print(index)  # tensor([[1, 0, 3, 2]])
print(index.t())
"""
tensor([[1],
        [0],
        [3],
        [2]])
"""

print(torch.gather(data, 0, index))  # tensor([[11, 15, 13,  8]])
print(torch.gather(data, 0, index.t()))
"""
tensor([[11],
        [ 5],
        [10],
        [ 1]])
"""
print(torch.gather(data, 1, index))  # tensor([[15,  5,  4,  6]])
print(torch.gather(data, 1, index.t()))
"""
tensor([[15],
        [11],
        [ 8],
        [13]])
"""

• torch.gather(data, 0, index)对0维(行)进行

    tensor([[1, 0, 3, 2]])
          第 0  1  2  3 列
  

  索引：
  [1][0] == 11
  [0][1] == 15
  [3][2] == 13
  [2][3] == 8
  加粗的是行索引的值1, 0, 3, 2，没加粗的为什么是0, 1, 2, 3呢？因为索引1是第0列，索引0是第1列，索引3是第2列，索引2是第3列。

• torch.gather(data, 0, index.t())对0维(行)进行

    tensor([[1],   第0列
            [0],   第0列
            [3],   第0列
            [2]])  第0列
  

  索引：
  [1][0] == 11
  [0][0] == 5
  [3][0] == 10
  [2][0] == 1
  加粗的是行索引的值1, 0, 3, 2，没加粗的都是0，因为索引1, 0, 3, 2都是第0列。

• torch.gather(data, 1, index.t())对1维(列)进行

    tensor([[1],   第0行
            [0],   第1行
            [3],   第2行
            [2]])  第3行
  

  索引：
  [0][1] == 15
  [1][0] == 11
  [2][3] == 8
  [3][2] == 13
  加粗的是列索引的值1, 0, 3, 2，没加粗的0, 1, 2, 3则是因为索引1是第0行，索引0是第1行，索引3是第2行，索引2是第3行。

.take/.take_along_dim

import torch

torch.manual_seed(1)

data = torch.randperm(16).view(4,4)
print(data)
"""
tensor([[ 5, 15,  6,  4],
        [11,  2,  7, 12],
        [ 1,  0,  9,  8],
        [10,  3, 13, 14]])
"""

index1 = torch.tensor([0, 5, 10, 15])
index2 = torch.tensor([[1, 0, 3, 2]])

print(torch.take(data, index1))  # tensor([ 5,  2,  9, 14])

print(torch.take_along_dim(data, index1))  # tensor([ 5,  2,  9, 14])
print(torch.take_along_dim(data, index2, dim=0))  # tensor([[11, 15, 13,  8]])
print(torch.take_along_dim(data, index2.t(), dim=0))  # 这里与torch.gather不同
"""
tensor([[11,  2,  7, 12],
        [ 5, 15,  6,  4],
        [10,  3, 13, 14],
        [ 1,  0,  9,  8]])
"""
print(torch.take_along_dim(data, index2, dim=1))  # 这里与torch.gather不同
"""
tensor([[15,  5,  4,  6],
        [ 2, 11, 12,  7],
        [ 0,  1,  8,  9],
        [ 3, 10, 14, 13]])
"""
print(torch.take_along_dim(data, index2.t(), dim=1))
"""
tensor([[15],
        [11],
        [ 8],
        [13]])
"""

torch.take_along_dim()当dim=None时，等价于torch.take()，先把张量打平转成1维在根据索引获取元素。

当dim不等于None时，则与data.gather()相似。

.argwhere/.nonzero

import torch

data1 = torch.tensor([1, 0, 1, 0, 1, 0])
data2 = torch.tensor([[1, 0, 1], [0, 1, 0]])
print(data1)  # tensor([1, 0, 1, 0, 1, 0])
print(data2)
"""
tensor([[1, 0, 1],
        [0, 1, 0]])
"""

print(torch.argwhere(data1))
"""
tensor([[0],
        [2],
        [4]])
"""
print(torch.argwhere(data2))
"""
tensor([[0, 0],
        [0, 2],
        [1, 1]])
"""

print(torch.nonzero(data1))
"""
tensor([[0],
        [2],
        [4]])
"""
print(torch.nonzero(data1, as_tuple=True))  # (tensor([0, 2, 4]),)
print(torch.nonzero(data2))
"""
tensor([[0, 0],
        [0, 2],
        [1, 1]])
"""
print(torch.nonzero(data2, as_tuple=True))  # (tensor([0, 0, 1]), tensor([0, 2, 1]))

torch.argwhere()和torch.nonzero()都是返回非0元素的索引。

当torch.nonzero()参数as_tuple=False时，效果与torch.argwhere()相同。

.where

import torch

torch.manual_seed(1)

condition = torch.randn(3, 5) > 0
print(condition)
"""
tensor([[ True,  True,  True,  True, False],
        [False, False,  True, False, False],
        [False, False, False, False,  True]])
"""

data1 = torch.ones(3, 5)
print(data1)
"""
tensor([[1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.]])
"""

data2 = torch.full_like(data1, 100)
print(data2)
"""
tensor([[100., 100., 100., 100., 100.],
        [100., 100., 100., 100., 100.],
        [100., 100., 100., 100., 100.]])
"""

# 满足条件返回 data1, 不满足条件返回 data2
print(torch.where(condition, data1, data2))
"""
tensor([[  1.,   1.,   1.,   1., 100.],
        [100., 100.,   1., 100., 100.],
        [100., 100., 100., 100.,   1.]])
"""

.unravel_index

import torch

print(torch.arange(9).view(3, 3))
"""
tensor([[0, 1, 2],
        [3, 4, 5],
        [6, 7, 8]])
"""
print(torch.arange(9, 18).view(3, 3))
"""
tensor([[ 9, 10, 11],
        [12, 13, 14],
        [15, 16, 17]])
"""

print(torch.unravel_index(torch.tensor(2), shape=(3, 3)))  # (tensor(0), tensor(2))
print(torch.unravel_index(torch.tensor([4, 6, 9, 17]), shape=(3, 3)))  # (tensor([1, 2, 0, 2]), tensor([1, 0, 0, 2]))

索引2在shape为(3, 3)的0行2列。
索引9可以理解为9 / prod(shape) == 9 % 9 == 0，所以在0行0列。
索引17可以理解为17 / prod(shape) == 17 % 9 == 8，所以在2行2列。

维度变换

.t/.transpose/.movedim/.permute

import torch

data = torch.arange(9)
print(data)  # tensor([0, 1, 2, 3, 4, 5, 6, 7, 8])
print(torch.t(data))  # tensor([0, 1, 2, 3, 4, 5, 6, 7, 8])

data = torch.arange(9).view(3, 3)
print(data)
"""
tensor([[0, 1, 2],
        [3, 4, 5],
        [6, 7, 8]])
"""
print(torch.t(data))
"""
tensor([[0, 3, 6],
        [1, 4, 7],
        [2, 5, 8]])
"""

# 4 张 3 通道，高 1080，宽 1920 的图片
images = torch.rand(4, 3, 1080, 1920)
# BCHW
print(images.shape)  # torch.Size([4, 3, 1080, 1920])

# BCHW --> BWHC --> BHWC
print(images.transpose(1, 3).transpose(1, 2).shape)  # torch.Size([4, 1080, 1920, 3])

# BCHW --> BHWC
print(images.movedim(1, 3).shape)  # torch.Size([4, 1080, 1920, 3])

# BCHW --> BHWC
print(images.permute(0, 2, 3, 1).shape)  # torch.Size([4, 1080, 1920, 3])

torch.t(input)只能处理维度小于等于2的，否则会报错。当维度是0或1维时，返回相同的结果，当维度为2时，等价于torch.transpose(input, 0, 1)

torch.transpose(input, dim0, dim1)一次只能操作两个维度，对调两个维度的位置。

torch.movedim(input, source, destination)将source维度移动到destination维度。

torch.permute(input, dims)一次可以操作多个维度，dims指定所有维度的顺序。在多维度操作上使用torch.permute()更直观。

torch.swapaxes()是torch.transpose()的别名。
torch.swapdims()是torch.transpose()的别名。
torch.moveaxis()是torch.movedim()的别名。

.view/.reshape

import torch

# 4 张 3 通道，高 28，宽 28 的图片
images = torch.rand(4, 3, 28, 28)
print(images.shape)  # torch.Size([4, 3, 28, 28])

print(images.view(4, 3 * 28 * 28).shape)  # torch.Size([4, 2352])
# -1 会自动计算出数值
print(images.view(4, -1).shape)  # torch.Size([4, 2352])
# print(images.transpose(1, 3).transpose(1, 2).view(4, -1).shape)  # 报错
print(images.transpose(1, 3).transpose(1, 2).contiguous().view(4, -1).shape)  # torch.Size([4, 2352])
# print(images.permute(0, 2, 3, 1).view(4, -1).shape)  # 报错
print(images.permute(0, 2, 3, 1).contiguous().view(4, -1).shape)  # torch.Size([4, 2352])

print(images.reshape(4, 3 * 28 * 28).shape)  # torch.Size([4, 2352])
# -1 会自动计算出数值
print(images.reshape(4, -1).shape)  # torch.Size([4, 2352])
print(images.transpose(1, 3).transpose(1, 2).reshape(4, -1).shape)  # torch.Size([4, 2352])
print(images.permute(0, 2, 3, 1).reshape(4, -1).shape)  # torch.Size([4, 2352])

view()和reshape()都可以改变Tensor的维度，区别是：

view()只能对满足连续性的张量进行转换，当对不满足连续性的张量进行操作时会报RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.错误。transpose()和permute()会改变张量连续性，使用view()前需要先执行contiguous()。

reshape()则没有上述要求，可以直接使用，无需先执行contiguous()。

.squeeze/.unsqueeze

import torch

data = torch.rand(2, 1, 2, 1)
print(data.shape)  # torch.Size([2, 1, 2, 1])

# dim=None，会把能挤压的维度挤压
print(torch.squeeze(data).shape)  # torch.Size([2, 2])
# 第一维度是 2 没法挤压，所以保持不变
print(torch.squeeze(data, dim=0).shape)  # torch.Size([2, 1, 2, 1])
# 第二维度挤压
print(torch.squeeze(data, dim=1).shape)  # torch.Size([2, 2, 1])

# 在第一维增加维度
print(torch.unsqueeze(data, dim=0).shape)  # torch.Size([1, 2, 1, 2, 1])
# 在最后一维增加维度
print(torch.unsqueeze(data, dim=-1).shape)  # torch.Size([2, 1, 2, 1, 1])

.expand/.repeat

import torch

# 4 张单通道，高 28，宽 28 的图片
images = torch.rand(4, 1, 28, 28)
data = torch.empty(4, 1, 1, 1)

# 第二维度扩展到 3 通道，-1 表示不扩展保持原样
print(images.expand(-1, 3, -1, -1).shape)  # torch.Size([4, 3, 28, 28])
print(images.expand(4, 3, 28, 28).shape)  # torch.Size([4, 3, 28, 28])
# 扩展到和 images 相同的 shape
print(data.expand_as(images).shape)  # torch.Size([4, 1, 28, 28])

# 各维度重复指定次数的数据
print(images.repeat(1, 3, 1, 1).shape)  # torch.Size([4, 3, 28, 28])

.tile

import torch

data = torch.arange(9).view(3, 3)
print(data)
"""
tensor([[0, 1, 2],
        [3, 4, 5],
        [6, 7, 8]])
"""

# 如果 dims 指定的维度小于 input 的维度，会扩展到相同的维度，比如 (2,) 会扩展到 (1, 2)
print(torch.tile(data, dims=(2,)))
"""
tensor([[0, 1, 2, 0, 1, 2],
        [3, 4, 5, 3, 4, 5],
        [6, 7, 8, 6, 7, 8]])
"""
print(torch.tile(data, dims=(1, 2)))
"""
tensor([[0, 1, 2, 0, 1, 2],
        [3, 4, 5, 3, 4, 5],
        [6, 7, 8, 6, 7, 8]])
"""
print(torch.tile(data, dims=(3, 2)))
"""
tensor([[0, 1, 2, 0, 1, 2],
        [3, 4, 5, 3, 4, 5],
        [6, 7, 8, 6, 7, 8],
        [0, 1, 2, 0, 1, 2],
        [3, 4, 5, 3, 4, 5],
        [6, 7, 8, 6, 7, 8],
        [0, 1, 2, 0, 1, 2],
        [3, 4, 5, 3, 4, 5],
        [6, 7, 8, 6, 7, 8]])
"""
# 如果 dims 指定的维度大于 input 的维度，input 会扩展到相同的维度，比如 input 是 (3, 3) 会扩展到 (1, 3, 3)
print(torch.tile(data, dims=(1, 3, 2)))
"""
tensor([[[0, 1, 2, 0, 1, 2],
         [3, 4, 5, 3, 4, 5],
         [6, 7, 8, 6, 7, 8],
         [0, 1, 2, 0, 1, 2],
         [3, 4, 5, 3, 4, 5],
         [6, 7, 8, 6, 7, 8],
         [0, 1, 2, 0, 1, 2],
         [3, 4, 5, 3, 4, 5],
         [6, 7, 8, 6, 7, 8]]])
"""

合并与拆分

.cat

import torch

torch.manual_seed(1)

data1 = torch.rand(2, 3)
print(data1)
"""
tensor([[0.7576, 0.2793, 0.4031],
        [0.7347, 0.0293, 0.7999]])
"""
data2 = torch.rand(2, 3)
print(data1)
"""
tensor([[0.7576, 0.2793, 0.4031],
        [0.7347, 0.0293, 0.7999]])
"""

# 行合并
print(torch.cat([data1, data1]))
"""
tensor([[0.7576, 0.2793, 0.4031],
        [0.7347, 0.0293, 0.7999],
        [0.7576, 0.2793, 0.4031],
        [0.7347, 0.0293, 0.7999]])
"""
# 列合并
print(torch.cat([data1, data1], dim=1))
"""
tensor([[0.7576, 0.2793, 0.4031, 0.7576, 0.2793, 0.4031],
        [0.7347, 0.0293, 0.7999, 0.7347, 0.0293, 0.7999]])
"""

torch.concat()是torch.cat()的别名。
torch.concatenate()是torch.cat()的别名。

.stack/.hstack/.vstack/.column_stack/.dstack

import torch

data = torch.tensor(9)
data1 = torch.arange(9)
data2 = torch.arange(9).view(3, 3)
data3 = torch.arange(9).view(1, 3, 3)
data4 = torch.arange(9).view(1, 1, 3, 3)
data5 = torch.arange(9).view(1, 1, 1, 3, 3)

print(data.shape)  # torch.Size([])
print(data1.shape)  # torch.Size([9])
print(data2.shape)  # torch.Size([3, 3])
print(data3.shape)  # torch.Size([1, 3, 3])
print(data4.shape)  # torch.Size([1, 1, 3, 3])
print(data5.shape)  # torch.Size([1, 1, 1, 3, 3])

# dim 默认为 0，在第一维度堆叠
print(torch.stack([data, data]).shape)  # torch.Size([2])
print(torch.stack([data1, data1]).shape)  # torch.Size([2, 9])
print(torch.stack([data2, data2]).shape)  # torch.Size([2, 3, 3])
print(torch.stack([data3, data3]).shape)  # torch.Size([2, 1, 3, 3])
print(torch.stack([data4, data4]).shape)  # torch.Size([2, 1, 1, 3, 3])
print(torch.stack([data5, data5]).shape)  # torch.Size([2, 1, 1, 1, 3, 3])

# dim=1，在第二维度堆叠
print(torch.stack([data1, data1], dim=1).shape)  # torch.Size([9, 2])
print(torch.stack([data2, data2], dim=1).shape)  # torch.Size([3, 2, 3])
print(torch.stack([data3, data3], dim=1).shape)  # torch.Size([1, 2, 3, 3])
print(torch.stack([data4, data4], dim=1).shape)  # torch.Size([1, 2, 1, 3, 3])
print(torch.stack([data5, data5], dim=1).shape)  # torch.Size([1, 2, 1, 1, 3, 3])

# 水平堆叠
print(torch.hstack([data, data]).shape)  # torch.Size([2])
print(torch.hstack([data1, data1]).shape)  # torch.Size([18])
print(torch.hstack([data2, data2]).shape)  # torch.Size([3, 6])
print(torch.hstack([data3, data3]).shape)  # torch.Size([1, 6, 3])
print(torch.hstack([data4, data4]).shape)  # torch.Size([1, 2, 3, 3])
print(torch.hstack([data5, data5]).shape)  # torch.Size([1, 2, 1, 3, 3])

# 垂直堆叠
print(torch.vstack([data, data]).shape)  # torch.Size([2, 1])
print(torch.vstack([data1, data1]).shape)  # torch.Size([2, 9])
print(torch.vstack([data2, data2]).shape)  # torch.Size([6, 3])
print(torch.vstack([data3, data3]).shape)  # torch.Size([2, 3, 3])
print(torch.vstack([data4, data4]).shape)  # torch.Size([2, 1, 3, 3])
print(torch.vstack([data5, data5]).shape)  # torch.Size([2, 1, 1, 3, 3])

# 列堆叠
print(torch.column_stack([data, data]).shape)  # torch.Size([1, 2])
print(torch.column_stack([data1, data1]).shape)  # torch.Size([9, 2])
print(torch.column_stack([data2, data2]).shape)  # torch.Size([3, 6])
print(torch.column_stack([data3, data3]).shape)  # torch.Size([1, 6, 3])
print(torch.column_stack([data4, data4]).shape)  # torch.Size([1, 2, 3, 3])
print(torch.column_stack([data5, data5]).shape)  # torch.Size([1, 2, 1, 3, 3])

# 深度堆叠，在第 3 维度堆叠
print(torch.dstack((data, data)).shape)  # torch.Size([1, 1, 2])
print(torch.dstack((data1, data1)).shape)  # torch.Size([1, 9, 2])
print(torch.dstack((data2, data2)).shape)  # torch.Size([3, 3, 2])
print(torch.dstack((data3, data3)).shape)  # torch.Size([1, 3, 6])
print(torch.dstack((data4, data4)).shape)  # torch.Size([1, 1, 6, 3])
print(torch.dstack((data5, data5)).shape)  # torch.Size([1, 1, 2, 3, 3])

torch.hstack)按水平方向（列方向）依次堆叠张量。

torch.vstack)按垂直方向（行方向）依次堆叠张量。

torch.column_stack()除了张量是0维和1维外，等价于torch.hstack()。当张量t是0维或1维时，先reshape重塑为(t.numel(), 1)再水平堆叠。

torch.row_stack()是torch.vstack()的别名。

torch.dstack()在第3维度进行堆叠。

torch.stack()是一个更通用的函数，通过dim指定在任意维度进行堆叠，dim默认等于0。它总是增加一个新的维度来堆叠张量。

.chunk

import torch

data = torch.arange(20).view(4, 5)
print(data)
"""
tensor([[ 0,  1,  2,  3,  4],
        [ 5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14],
        [15, 16, 17, 18, 19]])
"""

# 按行分成 3 块（PS: 这里实际只分成两块，因为每块大小是 2，4 行只能分成两块）
print(torch.chunk(data, 3))
"""
(tensor([[0, 1, 2, 3, 4],
        [5, 6, 7, 8, 9]]), tensor([[10, 11, 12, 13, 14],
        [15, 16, 17, 18, 19]]))
"""
# 按列分成 2 块
print(torch.chunk(data, 2, dim=1))
"""
(tensor([[ 0,  1,  2],
        [ 5,  6,  7],
        [10, 11, 12],
        [15, 16, 17]]), tensor([[ 3,  4],
        [ 8,  9],
        [13, 14],
        [18, 19]]))
"""

.split/.hsplit/.vsplit/.dsplit/.tensor_split

import torch

data1 = torch.arange(9)
data2 = torch.arange(9).view(3, 3)
data3 = torch.arange(9).view(1, 3, 3)

# 按每两个一组分割
print(torch.split(data1, 2))  # (tensor([0, 1]), tensor([2, 3]), tensor([4, 5]), tensor([6, 7]), tensor([8]))
# 分割成 3 组，第一组 1 份，第二组 3 份，第三组 5 份
print(torch.split(data1, [1, 3, 5]))  # (tensor([0]), tensor([1, 2, 3]), tensor([4, 5, 6, 7, 8]))
# 按每两个一组分割
print(torch.split(data2, 2, dim=1))
"""
(tensor([[0, 1],
        [3, 4],
        [6, 7]]), tensor([[2],
        [5],
        [8]]))
"""
# 分割成 2 组，第一组 1 份，第二组 2 份
print(torch.split(data2, [1, 2], dim=1))
"""
(tensor([[0],
        [3],
        [6]]), tensor([[1, 2],
        [4, 5],
        [7, 8]]))
"""

# 将具有一维或多维的张量 input 水平拆分为多个张量
# 平均分成 3 段
print(torch.hsplit(data1, 3))  # (tensor([0, 1, 2]), tensor([3, 4, 5]), tensor([6, 7, 8]))
# 以索引 3、5、7 分成 4 段
print(torch.hsplit(data1, [3, 5, 7]))  # (tensor([0, 1, 2]), tensor([3, 4]), tensor([5, 6]), tensor([7, 8]))
# 平均分成 3 段
print(torch.hsplit(data2, 3))
"""
(tensor([[0],
        [3],
        [6]]), tensor([[1],
        [4],
        [7]]), tensor([[2],
        [5],
        [8]]))
"""
# 以索引 1 分成 2 段
print(torch.hsplit(data2, [1]))
"""
(tensor([[0],
        [3],
        [6]]), tensor([[1, 2],
        [4, 5],
        [7, 8]]))
"""

# 将具有二维或多维的张量 input 垂直拆分为多个张量
# 平均分成 3 段
print(torch.vsplit(data2, 3))  # (tensor([[0, 1, 2]]), tensor([[3, 4, 5]]), tensor([[6, 7, 8]]))
# 以索引 2 分成 2 段
print(torch.vsplit(data2, [2]))
"""
(tensor([[0, 1, 2],
        [3, 4, 5]]), tensor([[6, 7, 8]]))
"""

# 将具有三个或更多维度的张量 input 在深度方向上拆分为多个张量
# 平均分成 3 段
print(torch.dsplit(data3, 3))
"""
(tensor([[[0],
         [3],
         [6]]]), tensor([[[1],
         [4],
         [7]]]), tensor([[[2],
         [5],
         [8]]]))
"""
# 以索引 1、2 分成 3 段
print(torch.dsplit(data3, [1, 2]))
"""
(tensor([[[0],
         [3],
         [6]]]), tensor([[[1],
         [4],
         [7]]]), tensor([[[2],
         [5],
         [8]]]))
"""

print(torch.tensor_split(data1, 3))  # (tensor([0, 1, 2]), tensor([3, 4, 5]), tensor([6, 7, 8]))
print(torch.tensor_split(data2, 3, dim=1))
"""
(tensor([[0],
        [3],
        [6]]), tensor([[1],
        [4],
        [7]]), tensor([[2],
        [5],
        [8]]))
"""
print(torch.tensor_split(data3, 3, dim=2))
"""
(tensor([[0],
        [3],
        [6]]), tensor([[1],
        [4],
        [7]]), tensor([[2],
        [5],
        [8]]))
"""

torch.hsplit(input, indices_or_sections)当input是1维时，等价于torch.tensor_split(input, indices_or_sections, dim=0)，当input是大于等于2维时，等价于torch.tensor_split(input, indices_or_sections, dim=1)。indices_or_sections如果是数字，必须能被整除，否则会抛出异常。

torch.vsplit(input, indices_or_sections)等价于torch.tensor_split(input, indices_or_sections, dim=0)。input必须大于等于2维。indices_or_sections如果是数字，必须能被整除，否则会抛出异常。

torch.dsplit(input, indices_or_sections)等价于torch.tensor_split(input, indices_or_sections, dim=2)。input必须大于等于3维。indices_or_sections如果是数字，必须能被整除，否则会抛出异常。

.unbind

import torch

data = torch.arange(9).view(3, 3)
print(data)
"""
tensor([[0, 1, 2],
        [3, 4, 5],
        [6, 7, 8]])
"""

# 按行解绑
print(torch.unbind(data))  # (tensor([0, 1, 2]), tensor([3, 4, 5]), tensor([6, 7, 8]))
# 按列解绑
print(torch.unbind(data, dim=1))  # (tensor([0, 3, 6]), tensor([1, 4, 7]), tensor([2, 5, 8]))

逐点运算

.add/.sub/.mul/.div/.remainder/.fmod/.positive/.neg/.abs

import torch

torch.manual_seed(1)

data = torch.randint(-10, 10, (2, 5))
print(data)
"""
tensor([[-5,  9, -6, -2, -7],
        [ 3,  1, -9,  9,  2]])
"""

# 加
print(torch.add(data, 2))  # 等价于 data + 2
"""
tensor([[-3, 11, -4,  0, -5],
        [ 5,  3, -7, 11,  4]])
"""

# 减
print(torch.sub(data, 2))  # 等价于 data - 2
"""
tensor([[ -7,   7,  -8,  -4,  -9],
        [  1,  -1, -11,   7,   0]])
"""

# 乘
print(torch.mul(data, 2))  # 等价于 data * 2
"""
tensor([[-10,  18, -12,  -4, -14],
        [  6,   2, -18,  18,   4]])
"""

# 除
print(torch.div(data, 2))  # 等价于 data / 2
"""
tensor([[-2.5000,  4.5000, -3.0000, -1.0000, -3.5000],
        [ 1.5000,  0.5000, -4.5000,  4.5000,  1.0000]])
"""
# 除完向下取整
print(torch.div(data, 2, rounding_mode='floor'))  # 等价于 data // 2
"""
tensor([[-3,  4, -3, -1, -4],
        [ 1,  0, -5,  4,  1]])
"""
# 除完只保留整数
print(torch.div(data, 2, rounding_mode='trunc'))
"""
tensor([[-2,  4, -3, -1, -3],
        [ 1,  0, -4,  4,  1]])
"""

# 取余不保留负数
print(torch.remainder(data, 2))  # 等价于 data % 2
"""
tensor([[1, 1, 0, 0, 1],
        [1, 1, 1, 1, 0]])
"""

# 取余保留负数
print(torch.fmod(data, 2))  # 等价于 data % -2
"""
tensor([[-1,  1,  0,  0, -1],
        [ 1,  1, -1,  1,  0]])
"""

# 加号
print(torch.positive(data))  # 等价于 +data
"""
tensor([[-5,  9, -6, -2, -7],
        [ 3,  1, -9,  9,  2]])
"""

# 减号
print(torch.neg(data))  # 等价于 -data
"""
tensor([[ 5, -9,  6,  2,  7],
        [-3, -1,  9, -9, -2]])
"""

# 绝对值
print(torch.abs(data))
"""
tensor([[5, 9, 6, 2, 7],
        [3, 1, 9, 9, 2]])
"""

torch.subtract()是torch.sub()的别名。
torch.multiply()是torch.mul()的别名。
torch.divide()是torch.div()的别名。
torch.true_divide()是torch.div()在rounding_mode=None时的别名。
torch.negative()是torch.neg()的别名。
torch.absolute()是torch.abs()的别名。

在PyTorch 1.13后（含1.13），可以认为torch.floor_divide()是torch.div()在rounding_mode='floor'时的别名，效果是一样的。

.pow/.square/.sqrt/.rsqrt/.reciprocal

import torch

data = torch.full((2, 2), 9)
print(data)
"""
tensor([[9, 9],
        [9, 9]])
"""

# 次方
print(torch.pow(data, 3))  # 等价于 data ** 3
"""
tensor([[729, 729],
        [729, 729]])
"""

# 平方
print(torch.square(data))  # 等价于 data ** 2
"""
tensor([[81, 81],
        [81, 81]])
"""

# 开方
print(torch.sqrt(data))  # 等价于 data ** 0.5
"""
tensor([[3., 3.],
        [3., 3.]])
"""

# 开方倒数
print(torch.rsqrt(data))  # 等价于 data ** -0.5
"""
tensor([[0.3333, 0.3333],
        [0.3333, 0.3333]])
"""

# 倒数
print(torch.reciprocal(data))  # 等价于 data ** -1
"""
tensor([[0.1111, 0.1111],
        [0.1111, 0.1111]])
"""

.exp/.log/.log2/.log10/.log1p

import torch

data1 = torch.exp(torch.tensor([0, 1, 2, 3, 4]))
print(data1)  # tensor([ 1.0000,  2.7183,  7.3891, 20.0855, 54.5981])

data2 = torch.exp(torch.ones(2, 2))
print(data2)
"""
tensor([[2.7183, 2.7183],
        [2.7183, 2.7183]])
"""

# 以 e 为底
print(torch.log(data1))  # tensor([0.0000, 1.0000, 2.0000, 3.0000, 4.0000])
print(torch.log(data2))
"""
tensor([[1.0000, 1.0000],
        [1.0000, 1.0000]])
"""

data3 = torch.tensor([1, 2, 4, 8, 16])
# 以 2 为底
print(torch.log2(data3))  # tensor([0., 1., 2., 3., 4.])

data4 = torch.tensor([1, 10, 100, 1000, 10000])
# 以 10 为底
print(torch.log10(data4))  # tensor([0., 1., 2., 3., 4.])

# 以 e 为底，input + 1 做为真数
print(torch.log1p(data1 - 1))  # tensor([0.0000, 1.0000, 2.0000, 3.0000, 4.0000])
print(torch.log1p(data2 - 1))
"""
tensor([[1.0000, 1.0000],
        [1.0000, 1.0000]])
"""

torch.log1p()计算input + 1的自然对数：$y_i = ln(x_i+1)$。

.sin/.cos/.tan/.asin/.acos/.atan/.atan2/.sinh/.cosh/.tanh

import torch

data = torch.deg2rad(torch.arange(0, 361, 30))  # 角度转弧度
print(data)

# 正弦
print(torch.sin(data))
# 余弦
print(torch.cos(data))
# 正切
print(torch.tan(data))

# 反正弦
print(torch.asin(data))
# 反余弦
print(torch.acos(data))
# 反正切
print(torch.atan(data))
print(torch.atan2(data, data))

# 双曲正弦
print(torch.sinh(data))
# 双曲余弦
print(torch.cosh(data))
# 双曲正切
print(torch.tanh(data))

torch.arcsin()是torch.asin()的别名。
torch.arccos()是torch.acos()的别名。
torch.arctan()是torch.atan()的别名。
torch.arctan2()是torch.atan2()的别名。
torch.arcsinh()是torch.asinh()的别名。
torch.arccosh()是torch.acosh()的别名。
torch.arctanh()是torch.atanh()的别名。

.angle/.deg2rad

import numpy as np
import torch

# 复数转弧度，结果在（π, -π）之间
print(torch.angle(torch.tensor([[1, 1 + 1j, 1j, -1 + 1j, -1, -1 -1j, -1j, 1 -1j]])))  # tensor([[ 0.0000,  0.7854,  1.5708,  2.3562,  3.1416, -2.3562, -1.5708, -0.7854]])
# 弧度转角度
print(torch.angle(torch.tensor([[1, 1 + 1j, 1j, -1 + 1j, -1, -1 -1j, -1j, 1 -1j]])) * 180 / np.pi)  # tensor([[   0.0000,   45.0000,   90.0000,  135.0000,  180.0000, -135.0000, -90.0000,  -45.0000]])

# 角度转弧度
print(torch.deg2rad(torch.tensor([[0, 45, 90, 135, 180, 360], [-360, -180, -135, -90, -45, 0]])))
"""
tensor([[ 0.0000,  0.7854,  1.5708,  2.3562,  3.1416,  6.2832],
        [-6.2832, -3.1416, -2.3562, -1.5708, -0.7854,  0.0000]])
"""

.bitwise_and/.bitwise_or/.bitwise_not/.bitwise_xor/.bitwise_left_shift/.bitwise_right_shift

import torch

# 按位与
print(torch.bitwise_and(torch.tensor([-1, -2, 3], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)))  # tensor([1, 0, 3], dtype=torch.int8)
print(torch.bitwise_and(torch.tensor([True, True, False]), torch.tensor([False, True, False])))  # tensor([False,  True, False])

# 按位或
print(torch.bitwise_or(torch.tensor([-1, -2, 3], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)))  # tensor([-1, -2,  3], dtype=torch.int8)
print(torch.bitwise_or(torch.tensor([True, True, False]), torch.tensor([False, True, False])))  # tensor([ True,  True, False])

# 按位非
print(torch.bitwise_not(torch.tensor([-1, -2, 3], dtype=torch.int8)))  # tensor([ 0,  1, -4], dtype=torch.int8)
print(torch.bitwise_not(torch.tensor([True, True, False])))  # tensor([False, False,  True])

# 按位异或
print(torch.bitwise_xor(torch.tensor([-1, -2, 3], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)))  # tensor([-2, -2,  0], dtype=torch.int8)
print(torch.bitwise_xor(torch.tensor([True, True, False]), torch.tensor([False, True, False])))  # tensor([ True, False, False])

# 左移
print(torch.bitwise_left_shift(torch.tensor([-1, -2, 3], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)))  # tensor([-2, -2, 24], dtype=torch.int8)
# 右移
print(torch.bitwise_right_shift(torch.tensor([-1, -2, 3], dtype=torch.int8), torch.tensor([1, 0, 3], dtype=torch.int8)))  # tensor([-1, -2,  0], dtype=torch.int8)

.floor/.ceil/.round

import torch

torch.manual_seed(1)

data = torch.rand(1, 5) * 10
print(data)  # tensor([[7.5763, 2.7931, 4.0307, 7.3468, 0.2928]])

# 向下取整
print(torch.floor(data))  # tensor([[7., 2., 4., 7., 0.]])
# 向上取整
print(torch.ceil(data))  # tensor([[8., 3., 5., 8., 1.]])
# 四舍五入
print(torch.round(data))  # tensor([[8., 3., 4., 7., 0.]])

torch.floor()向下取整，torch.ceil()向上取整，torch.round()四舍五入。

.trunc/.frac

import torch

torch.manual_seed(1)

data = torch.rand(1, 5) * 10
print(data)  # tensor([[7.5763, 2.7931, 4.0307, 7.3468, 0.2928]])

# 取整数
print(torch.trunc(data))  # tensor([[7., 2., 4., 7., 0.]])
# 取小数
print(torch.frac(data))  # tensor([[0.5763, 0.7931, 0.0307, 0.3468, 0.2928]])

torch.trunc()取整数，torch.frac()取小数。

torch.fix()是torch.trunc()的别名。

.clamp/.clamp_min/.clamp_max

import torch

torch.manual_seed(1)

data = torch.rand((2, 5)) * 20
print(data)
"""
tensor([[15.1526,  5.5862,  8.0614, 14.6937,  0.5856],
        [15.9972,  7.9427, 15.0874, 11.3902,  8.7756]])
"""

# 小于 10 的数变为 10
print(torch.clamp(data, 10))
"""
tensor([[15.1526, 10.0000, 10.0000, 14.6937, 10.0000],
        [15.9972, 10.0000, 15.0874, 11.3902, 10.0000]])
"""

# 大于 15 的数变为 15
print(torch.clamp(data, max=15))
"""
tensor([[15.0000,  5.5862,  8.0614, 14.6937,  0.5856],
        [15.0000,  7.9427, 15.0000, 11.3902,  8.7756]])
"""

# 小于 10 的数变为 10，大于 15 的数变为 15
print(torch.clamp(data, 10, 15))
"""
tensor([[15.0000, 10.0000, 10.0000, 14.6937, 10.0000],
        [15.0000, 10.0000, 15.0000, 11.3902, 10.0000]])
"""

# 小于 10 的数变为 10
print(torch.clamp_min(data, 10))
"""
tensor([[15.1526, 10.0000, 10.0000, 14.6937, 10.0000],
        [15.9972, 10.0000, 15.0874, 11.3902, 10.0000]])
"""

# 大于 15 的数变为 15
print(torch.clamp_max(data, 15))
"""
tensor([[15.0000,  5.5862,  8.0614, 14.6937,  0.5856],
        [15.0000,  7.9427, 15.0000, 11.3902,  8.7756]])
"""

torch.clamp()用于将输入的张量夹紧到区间[min, max]。

torch.clip()是torch.clamp()的别名。

.dot/.mm/.bmm/.matmul

import torch

# 一维矩阵相乘
print(torch.dot(torch.tensor([1, 2]), torch.tensor([3, 4])))  # tensor(11)

# 二维矩阵相乘
print(torch.mm(torch.full((2, 3), 2), torch.full((3, 2), 3)))
"""
tensor([[18, 18],
        [18, 18]])
"""

# 三维矩阵相乘
print(torch.bmm(torch.full((2, 2, 3), 2), torch.full((2, 3, 2), 3)))
"""
tensor([[[18, 18],
         [18, 18]],

        [[18, 18],
         [18, 18]]])
"""

data1 = torch.rand(4, 3, 800, 600)
print(data1.shape)  # torch.Size([4, 3, 800, 600])
data2 = torch.rand(4, 3, 600, 400)
print(data2.shape)  # torch.Size([4, 3, 600, 400])
# data1 @ data2 等价于 data1.matmul(data2)
print(data1.matmul(data2).shape)  # torch.Size([4, 3, 800, 400])

dot()只支持一维矩阵相乘，mm()只支持二维矩阵相乘，bmm()只支持三维矩阵相乘，matmul()支持任意维度矩阵相乘。

比较运算

.eq/.ne/.gt/.ge/.lt/.le/.equal

import torch

torch.manual_seed(1)

data = torch.randperm(10).view(2, 5)
print(data)
"""
tensor([[5, 6, 1, 2, 0],
        [8, 9, 3, 7, 4]])
"""

# 等于
print(torch.eq(data, 5))  # 等价于 data == 5
"""
tensor([[ True, False, False, False, False],
        [False, False, False, False, False]])
"""

# 不等于
print(torch.ne(data, 5))  # 等价于 data != 5
"""
tensor([[False,  True,  True,  True,  True],
        [ True,  True,  True,  True,  True]])
"""

# 大于
print(torch.gt(data, 5))  # 等价于 data > 5
"""
tensor([[False,  True, False, False, False],
        [ True,  True, False,  True, False]])
"""

# 大于等于
print(torch.ge(data, 5))  # 等价于 data >= 5
"""
tensor([[ True,  True, False, False, False],
        [ True,  True, False,  True, False]])
"""

# 小于
print(torch.lt(data, 5))  # 等价于 data < 5
"""
tensor([[False, False,  True,  True,  True],
        [False, False,  True, False,  True]])
"""

# 小于等于
print(torch.le(data, 5))  # 等价于 data <= 5
"""
tensor([[ True, False,  True,  True,  True],
        [False, False,  True, False,  True]])
"""

# 比较两个 tensor 是否相同
print(torch.equal(torch.tensor([1, 2]), torch.tensor([1, 2])))  # True
print(torch.equal(torch.tensor([1, 2]), torch.tensor([2, 1])))  # False

torch.not_equal()是torch.ne()的别名。
torch.greater()是torch.gt()的别名。
torch.greater_equal()是torch.ge()的别名。
torch.less()是torch.lt()的别名。
torch.less_equal()是torch.le()的别名。

.isfinite/.isinf/.isposinf/.isneginf/.isnan/.isreal/.isin

import torch

data = torch.tensor([1, float('inf'), torch.inf, float('-inf'), -torch.inf, float('nan'), torch.nan, True, False])
# 是否不是无穷大
print(torch.isfinite(data))  # tensor([ True, False, False, False, False, False, False,  True,  True])
# 是否是无穷大
print(torch.isinf(data))  # tensor([False,  True,  True,  True,  True, False, False, False, False])
# 是否是正无穷大
print(torch.isposinf(data))  # tensor([False,  True,  True, False, False, False, False, False, False])
# 是否是负无穷大
print(torch.isneginf(data))  # tensor([False, False, False,  True,  True, False, False, False, False])
# 是否是 NaN
print(torch.isnan(data))  # tensor([False, False, False, False, False,  True,  True, False, False])
# 是否是实数
print(torch.isreal(torch.tensor([1, 1+1j, 2+0j, float('nan'), True, False])))  # tensor([ True, False,  True,  True,  True,  True])

elements = [[1, 2], [3, 4]]
test_elements = [2, 3]
# 对 elements 的每个元素进行判断是否在 test_elements 中
print(torch.isin(torch.tensor(elements), torch.tensor(test_elements)))
"""
tensor([[False,  True],
        [ True, False]])
"""

.isclose/.allclose

import torch

# 是否接近
print(torch.isclose(torch.tensor((1., 2, 3)), torch.tensor((1 + 1e-10, 3, 4))))  # tensor([ True, False, False])
print(torch.isclose(torch.tensor((float('inf'), 4)), torch.tensor((float('inf'), 6)), rtol=.5))  # tensor([True, True])

# 是否都接近
print(torch.allclose(torch.tensor([10000., 1e-07]), torch.tensor([10000.1, 1e-08])))  # False
print(torch.allclose(torch.tensor([10000., 1e-08]), torch.tensor([10000.1, 1e-09])))  # True
print(torch.allclose(torch.tensor([1.0, float('nan')]), torch.tensor([1.0, float('nan')])))  # False
print(torch.allclose(torch.tensor([1.0, float('nan')]), torch.tensor([1.0, float('nan')]), equal_nan=True))  # True

torch.isclose(input, other, rtol=1e-05, atol=1e-08, equal_nan=False)

torch.allclose(input, other, rtol=1e-05, atol=1e-08, equal_nan=False)

校验是否接近公式：$|input - other| \le atol + rtol \times |other|$

.maximum/.minimum/.fmax/.fmin

import torch

data1 = torch.tensor([9.7, float('nan'), 3.1, torch.nan, 11.1])
data2 = torch.tensor([-2.2, 0.5, torch.nan, float('nan'), 7.8])

print(torch.maximum(data1, data2))  # tensor([ 9.7000,     nan,     nan,     nan, 11.1000])
print(torch.minimum(data1, data2))  # tensor([-2.2000,     nan,     nan,     nan,  7.8000])

print(torch.fmax(data1, data2))  # tensor([ 9.7000,  0.5000,  3.1000,     nan, 11.1000])
print(torch.fmin(data1, data2))  # tensor([-2.2000,  0.5000,  3.1000,     nan,  7.8000])

torch.maximum()和torch.minimum()用于比较两数的大小，不支持比较NaN。

torch.fmax()和torch.fmin()用于比较两数的大小，支持比较NaN。

.sort/.msort/.argsort

import torch

torch.manual_seed(1)

data = torch.rand(2, 5)
print(data)
"""
tensor([[0.7576, 0.2793, 0.4031, 0.7347, 0.0293],
        [0.7999, 0.3971, 0.7544, 0.5695, 0.4388]])
"""

### dim 默认等于 -1，按列进行排序
print(torch.sort(data))
"""
torch.return_types.sort(
values=tensor([[0.0293, 0.2793, 0.4031, 0.7347, 0.7576],
        [0.3971, 0.4388, 0.5695, 0.7544, 0.7999]]),
indices=tensor([[4, 1, 2, 3, 0],
        [1, 4, 3, 2, 0]]))
"""
# 按行进行排序
print(torch.sort(data, dim=0))
"""
torch.return_types.sort(
values=tensor([[0.7576, 0.2793, 0.4031, 0.5695, 0.0293],
        [0.7999, 0.3971, 0.7544, 0.7347, 0.4388]]),
indices=tensor([[0, 0, 0, 1, 0],
        [1, 1, 1, 0, 1]]))
"""
# descending=True 降序排序
print(torch.sort(data, descending=True))
"""
torch.return_types.sort(
values=tensor([[0.7576, 0.7347, 0.4031, 0.2793, 0.0293],
        [0.7999, 0.7544, 0.5695, 0.4388, 0.3971]]),
indices=tensor([[0, 3, 2, 1, 4],
        [0, 2, 3, 4, 1]]))
"""

# 沿第一维度进行排序
print(torch.msort(data))  # 等价于 torch.sort(data, dim=0)[0]
"""
tensor([[0.7576, 0.2793, 0.4031, 0.5695, 0.0293],
        [0.7999, 0.3971, 0.7544, 0.7347, 0.4388]])
"""

### dim 默认等于 -1，按列进行排序，返回索引
print(torch.argsort(data))
"""
tensor([[4, 1, 2, 3, 0],
        [1, 4, 3, 2, 0]])
"""
# 按行进行排序，返回索引
print(torch.argsort(data, dim=0))
"""
tensor([[0, 0, 0, 1, 0],
        [1, 1, 1, 0, 1]])
"""
# descending=True 降序排序，返回索引
print(torch.argsort(data, descending=True))
"""
tensor([[0, 3, 2, 1, 4],
        [0, 2, 3, 4, 1]])
"""

torch.sort()排序后返回排序结果和索引。

torch.msort(input))等价于torch.sort()对第一维度进行排序在取排序结果，torch.sort(input, dim=0)[0]，返回结果不包含索引。

torch.argsort()返回排序后的索引。

.topk/.kthvalue

import torch

torch.manual_seed(1)

data = torch.randperm(30, dtype=torch.double).view(3, 10)
print(data)
"""
tensor([[25.,  4.,  6.,  8., 23., 18., 17., 20., 19., 12.],
        [ 5., 14., 22.,  3., 27., 15.,  9., 13.,  7., 11.],
        [10.,  2., 24., 29., 21., 26., 28.,  1., 16.,  0.]],
       dtype=torch.float64)
"""

# 获取前 k 大的数据
print(torch.topk(data, 3))
"""
torch.return_types.topk(
values=tensor([[25., 23., 20.],
        [27., 22., 15.],
        [29., 28., 26.]], dtype=torch.float64),
indices=tensor([[0, 4, 7],
        [4, 2, 5],
        [3, 6, 5]]))
"""
# 获取前 k 小的数据
print(torch.topk(data, 3, largest=False))
"""
torch.return_types.topk(
values=tensor([[4., 6., 8.],
        [3., 5., 7.],
        [0., 1., 2.]], dtype=torch.float64),
indices=tensor([[1, 2, 3],
        [3, 0, 8],
        [9, 7, 1]]))
"""

# 获取第 k 小的数据
print(torch.kthvalue(data, 8))
"""
torch.return_types.kthvalue(
values=tensor([20., 15., 26.], dtype=torch.float64),
indices=tensor([7, 5, 5]))
"""
print(torch.kthvalue(data, 2, dim=0))
"""
torch.return_types.kthvalue(
values=tensor([10.,  4., 22.,  8., 23., 18., 17., 13., 16., 11.], dtype=torch.float64),
indices=tensor([2, 0, 1, 0, 0, 0, 0, 1, 2, 1]))
"""

torch.topk()获取前k大的数据，torch.kthvalue()获取第k小的数据。

归约运算

.max/.min/.amax/.amin/.aminmax/.argmax/.argmin

import torch

torch.manual_seed(1)

data = torch.rand(4, 4)
print(data)
"""
tensor([[0.7576, 0.2793, 0.4031, 0.7347],
        [0.0293, 0.7999, 0.3971, 0.7544],
        [0.5695, 0.4388, 0.6387, 0.5247],
        [0.6826, 0.3051, 0.4635, 0.4550]])
"""

# 最大值
print(torch.max(data))  # tensor(0.7999)
print(torch.max(data, dim=0))
"""
torch.return_types.max(
values=tensor([0.7576, 0.7999, 0.6387, 0.7544]),
indices=tensor([0, 1, 2, 1]))
"""

# 最小值
print(torch.min(data))  # tensor(0.0293)
print(torch.min(data, dim=1))
"""
torch.return_types.min(
values=tensor([0.2793, 0.0293, 0.4388, 0.3051]),
indices=tensor([1, 0, 1, 1]))
"""

# 最大值
print(torch.amax(data))  # tensor(0.7999)
print(torch.amax(data, dim=0))  # tensor([0.7576, 0.7999, 0.6387, 0.7544])

# 最小值
print(torch.amin(data))  # tensor(0.0293)
print(torch.amin(data, dim=1))  # tensor([0.2793, 0.0293, 0.4388, 0.3051])

# 最大值和最小值
print(torch.aminmax(data))
"""
torch.return_types.aminmax(
min=tensor(0.0293),
max=tensor(0.7999))
"""
print(torch.aminmax(data, dim=0))
"""
torch.return_types.aminmax(
min=tensor([0.0293, 0.2793, 0.3971, 0.4550]),
max=tensor([0.7576, 0.7999, 0.6387, 0.7544]))
"""

# 最大值索引
print(torch.argmax(data))  # tensor(5)
print(torch.argmax(data, dim=0))  # tensor([0, 1, 2, 1])

# 最小值索引
print(torch.argmin(data))  # tensor(4)
print(torch.argmin(data, dim=1))  # tensor([1, 0, 1, 1])

torch.max()和torch.min()指定了dim时，返回极值和索引。

torch.amax()和torch.amin()只返回极值，不返回索引。

torch.argmax()和torch.argmin()返回极值对应的索引。不指定dim时，返回的是打平后的索引。

.mean/.nanmean

import torch

data = torch.tensor([[torch.nan, 1, 2, 3, 4], [5, 6, 7, 8, 9], [10, 11, 12, 13, torch.nan]])
print(data)
"""
tensor([[nan,  1.,  2.,  3.,  4.],
        [ 5.,  6.,  7.,  8.,  9.],
        [10., 11., 12., 13., nan]])
"""

# 平均值
print(torch.mean(data))  # tensor(nan)
print(torch.mean(data, dim=0))  # tensor([nan, 6., 7., 8., nan])
print(torch.mean(data, dim=1))  # tensor([nan, 7., nan])

# 平均值（忽略 NaN 值）
print(torch.nanmean(data))  # tensor(7.)
print(torch.nanmean(data, dim=0))  # tensor([7.5000, 6.0000, 7.0000, 8.0000, 6.5000])
print(torch.nanmean(data, dim=1))  # tensor([ 2.5000,  7.0000, 11.5000])

.median/.nanmedian

import torch

data = torch.tensor([[torch.nan, 1, 2, 3, 4], [5, 6, 7, 8, 9], [10, 11, 12, 13, torch.nan]])
print(data)
"""
tensor([[nan,  1.,  2.,  3.,  4.],
        [ 5.,  6.,  7.,  8.,  9.],
        [10., 11., 12., 13., nan]])
"""

# 中位数
print(torch.median(data))  # tensor(nan)
print(torch.median(data, dim=0))
"""
torch.return_types.median(
values=tensor([nan, 6., 7., 8., nan]),
indices=tensor([0, 1, 1, 1, 2]))
"""
print(torch.median(data, dim=1))  # tensor([nan, 7., nan])
"""
torch.return_types.median(
values=tensor([nan, 7., nan]),
indices=tensor([0, 2, 4]))
"""

# 中位数（忽略 NaN 值）
print(torch.nanmedian(data))  # tensor(7.)
print(torch.nanmedian(data, dim=0))
"""
torch.return_types.nanmedian(
values=tensor([5., 6., 7., 8., 4.]),
indices=tensor([1, 1, 1, 1, 0]))
"""
print(torch.nanmedian(data, dim=1))
"""
torch.return_types.nanmedian(
values=tensor([ 2.,  7., 11.]),
indices=tensor([2, 2, 1]))
"""

.sum/.nansum

import torch

data = torch.tensor([[torch.nan, 1, 2, 3, 4], [5, 6, 7, 8, 9], [10, 11, 12, 13, torch.nan]])
print(data)
"""
tensor([[nan,  1.,  2.,  3.,  4.],
        [ 5.,  6.,  7.,  8.,  9.],
        [10., 11., 12., 13., nan]])
"""

# 累加
print(torch.sum(data))  # tensor(nan)
print(torch.sum(data, dim=0))  # tensor([nan, 18., 21., 24., nan])
print(torch.sum(data, dim=1))  # tensor([nan, 35., nan])

# 累加（忽略 NaN 值）
print(torch.nansum(data))  # tensor(91.)
print(torch.nansum(data, dim=0))  # tensor([15., 18., 21., 24., 13.])
print(torch.nansum(data, dim=1))  # tensor([10., 35., 46.])

.prod

import torch

data = torch.tensor([[torch.nan, 1, 2, 3, 4], [5, 6, 7, 8, 9], [10, 11, 12, 13, torch.nan]])
print(data)
"""
tensor([[nan,  1.,  2.,  3.,  4.],
        [ 5.,  6.,  7.,  8.,  9.],
        [10., 11., 12., 13., nan]])
"""

# 累乘
print(torch.prod(data))  # tensor(nan)
print(torch.prod(data, dim=0))  # tensor([ nan,  66., 168., 312.,  nan])
print(torch.prod(data, dim=1))  # tensor([   nan, 15120.,    nan])

.var/.var_mean/.std/.std_mean

import torch

data = torch.arange(1, 10, dtype=torch.double).view(3, 3)
print(data)
"""
tensor([[1., 2., 3.],
        [4., 5., 6.],
        [7., 8., 9.]], dtype=torch.float64)
"""

# 样本方差
print(torch.var(data))  # tensor(7.5000, dtype=torch.float64)
# 总体方差
print(torch.var(data, unbiased=False))  # tensor(6.6667, dtype=torch.float64)
# 同时计算方差和平均值
print(torch.var_mean(data))  # (tensor(7.5000, dtype=torch.float64), tensor(5., dtype=torch.float64))

# 样本标准差
print(torch.std(data))  # tensor(2.7386, dtype=torch.float64)
# 总体标准差
print(torch.std(data, unbiased=False))  # tensor(2.5820, dtype=torch.float64)
# 同时计算标准差和平均值
print(torch.std_mean(data))  # (tensor(2.7386, dtype=torch.float64), tensor(5., dtype=torch.float64))

总体方差：$\sigma^2 = \frac{1}{N}\sum_{i=1}^{N}(x_i - \mu)^2$

样本方差：$s^2 = \frac{1}{n-1}\sum_{i=1}^{n}(x_i - \bar{x})^2$

总体标准差：$\sigma = \sqrt{\frac{1}{N}\sum_{i=1}^{N}(x_i - \mu)^2}$

样本标准差：$s = \sqrt{\frac{1}{n-1}\sum_{i=1}^{n}(x_i - \bar{x})^2}$

.norm

import torch

data = torch.full((2, 2), 2.)
print(data)
"""
tensor([[2., 2.],
        [2., 2.]])
"""

# 一范数
print(torch.norm(data, p=1))  # tensor(8.)
# 二范数
print(torch.norm(data))  # tensor(4.)
# 三范数
print(torch.norm(data, p=3))  # tensor(3.1748)

一范数：$||x||_1 = \sum_{i=1}^N|x_i|$

二范数：$||x||_2 = \sqrt{\sum_{i=1}^Nx_i^2}$

p范数：$||x||_p = (\sum_{i=1}^N|x_i|^p)^\frac{1}{p}$

.dist

import torch

input = torch.full((2, 2), 4.)
print(input)
"""
tensor([[4., 4.],
        [4., 4.]])
"""
other = torch.full((2, 2), 2.)
print(other)
"""
tensor([[2., 2.],
        [2., 2.]])
"""

# 一范数
print(torch.dist(input, other, p=1))  # tensor(8.)
# 二范数
print(torch.dist(input, other))  # tensor(4.)
# 三范数
print(torch.dist(input, other, p=3))  # tensor(3.1748)

torch.dist()返回(input - other)的p范数。

.any/.all

import torch

print(torch.any(torch.tensor([])))  # tensor(False)
print(torch.any(torch.tensor([True, True])))  # tensor(True)
print(torch.any(torch.tensor([True, False])))  # tensor(True)
print(torch.any(torch.tensor([False, False])))  # tensor(False)

print(torch.all(torch.tensor([])))  # tensor(True)
print(torch.all(torch.tensor([True, True])))  # tensor(True)
print(torch.all(torch.tensor([True, False])))  # tensor(False)
print(torch.all(torch.tensor([False, False])))  # tensor(False)

序列化

.save

import io
import torch

# Save to file
x = torch.tensor([0, 1, 2, 3, 4])
torch.save(x, 'tensor.pt')

# Save to io.BytesIO buffer
buffer = io.BytesIO()
torch.save(x, buffer)

.load

import io
import torch

# Load a module with 'ascii' encoding for unpickling
torch.load('tensors.pt', encoding='ascii')
# Load all tensors onto the CPU
torch.load('tensors.pt', map_location=torch.device('cpu'))
# Map tensors from GPU 1 to GPU 0
torch.load('tensors.pt', map_location={'cuda:1': 'cuda:0'})

# Load from io.BytesIO buffer
with open('tensor.pt', 'rb') as f:
    buffer = io.BytesIO(f.read())
torch.load(buffer)