跳到主要内容

张量操作

张量（Tensor）是 PyTorch 中最核心的数据结构，它是多维数组的泛化形式。理解张量操作是学习 PyTorch 的第一步。

什么是张量？​

张量是一个多维数组，可以看作是标量、向量和矩阵的推广：

维度名称示例
0维标量（Scalar）5
1维向量（Vector）[1, 2, 3]
2维矩阵（Matrix）[[1, 2], [3, 4]]
3维3D张量RGB图像（高×宽×通道）
n维nD张量批量数据（批次×通道×高×宽）
张量维度示意：

标量 (0D)     向量 (1D)      矩阵 (2D)        3D张量
    5        [1,2,3]    [[1,2,3],     [[[1,2],
                            [4,5,6]]      [3,4]],
                                         [[5,6],
                                          [7,8]]]

创建张量​

从 Python 列表创建​

import torch

# 从列表创建
x = torch.tensor([1, 2, 3, 4, 5])
print(x)
# tensor([1, 2, 3, 4, 5])

# 创建二维张量（矩阵）
matrix = torch.tensor([[1, 2, 3], [4, 5, 6]])
print(matrix)
# tensor([[1, 2, 3],
#         [4, 5, 6]])

# 指定数据类型
x_float = torch.tensor([1, 2, 3], dtype=torch.float32)
print(x_float.dtype)
# torch.float32

使用工厂函数创建​

# 创建全零张量
zeros = torch.zeros(3, 4)
print(zeros)
# tensor([[0., 0., 0., 0.],
#         [0., 0., 0., 0.],
#         [0., 0., 0., 0.]])

# 创建全一张量
ones = torch.ones(2, 3)
print(ones)
# tensor([[1., 1., 1.],
#         [1., 1., 1.]])

# 创建单位矩阵
eye = torch.eye(3)
print(eye)
# tensor([[1., 0., 0.],
#         [0., 1., 0.],
#         [0., 0., 1.]])

# 创建未初始化张量（值不确定）
empty = torch.empty(2, 3)
print(empty)

# 创建填充特定值的张量
full = torch.full((2, 3), fill_value=7)
print(full)
# tensor([[7, 7, 7],
#         [7, 7, 7]])

创建随机张量​

# 均匀分布 [0, 1)
rand = torch.rand(2, 3)
print(rand)

# 标准正态分布
randn = torch.randn(2, 3)
print(randn)

# 随机整数
randint = torch.randint(low=0, high=10, size=(2, 3))
print(randint)

# 设置随机种子（保证可重复性）
torch.manual_seed(42)
rand1 = torch.rand(2, 2)
torch.manual_seed(42)
rand2 = torch.rand(2, 2)
print(torch.equal(rand1, rand2))  # True

创建序列张量​

# 类似 range
arange = torch.arange(start=0, end=10, step=2)
print(arange)
# tensor([0, 2, 4, 6, 8])

# 等间隔张量
linspace = torch.linspace(start=0, end=1, steps=5)
print(linspace)
# tensor([0.0000, 0.2500, 0.5000, 0.7500, 1.0000])

# 对数间隔
logspace = torch.logspace(start=0, end=2, steps=5)
print(logspace)
# tensor([1.0000, 3.1623, 10.0000, 31.6228, 100.0000])

从 NumPy 创建​

import numpy as np

# NumPy 数组转张量
np_array = np.array([1, 2, 3, 4])
tensor = torch.from_numpy(np_array)
print(tensor)
# tensor([1, 2, 3, 4], dtype=torch.int32)

# 张量转 NumPy 数组
tensor = torch.tensor([1, 2, 3, 4])
np_array = tensor.numpy()
print(np_array)
# [1 2 3 4]

基于现有张量创建​

x = torch.tensor([[1, 2, 3], [4, 5, 6]])

# 相同形状的全零张量
zeros_like = torch.zeros_like(x)
print(zeros_like)
# tensor([[0, 0, 0],
#         [0, 0, 0]])

# 相同形状的全一张量
ones_like = torch.ones_like(x)
print(ones_like)

# 相同形状的随机张量
rand_like = torch.rand_like(x, dtype=torch.float)
print(rand_like)

张量属性​

x = torch.randn(2, 3, 4)

# 形状
print(x.shape)        # torch.Size([2, 3, 4])
print(x.size())       # torch.Size([2, 3, 4])

# 维度数
print(x.dim())        # 3

# 元素总数
print(x.numel())      # 24

# 数据类型
print(x.dtype)        # torch.float32

# 存储设备
print(x.device)       # cpu 或 cuda:0

数据类型​

PyTorch 支持多种数据类型：

类型说明对应 NumPy
torch.float32 / torch.float32位浮点数float32
torch.float64 / torch.double64位浮点数float64
torch.float16 / torch.half16位浮点数float16
torch.int32 / torch.int32位整数int32
torch.int64 / torch.long64位整数int64
torch.int16 / torch.short16位整数int16
torch.int88位整数int8
torch.uint88位无符号整数uint8
torch.bool布尔类型bool
# 类型转换
x = torch.tensor([1, 2, 3])
print(x.dtype)  # torch.int64

x_float = x.float()  # 或 x.to(torch.float32)
print(x_float.dtype)  # torch.float32

x_long = x_float.long()  # 或 x.to(torch.int64)
print(x_long.dtype)  # torch.int64

张量运算​

算术运算​

x = torch.tensor([1.0, 2.0, 3.0])
y = torch.tensor([4.0, 5.0, 6.0])

# 加法
print(x + y)           # tensor([5., 7., 9.])
print(torch.add(x, y)) # tensor([5., 7., 9.])

# 减法
print(x - y)           # tensor([-3., -3., -3.])

# 乘法（逐元素）
print(x * y)           # tensor([4., 10., 18.])

# 除法
print(x / y)           # tensor([0.2500, 0.4000, 0.5000])

# 幂运算
print(x ** 2)          # tensor([1., 4., 9.])

# 开方
print(torch.sqrt(x))   # tensor([1.0000, 1.4142, 1.7321])

# 指数和对数
print(torch.exp(x))    # e^x
print(torch.log(x))    # ln(x)

矩阵运算​

A = torch.tensor([[1, 2], [3, 4]], dtype=torch.float)
B = torch.tensor([[5, 6], [7, 8]], dtype=torch.float)

# 矩阵乘法
print(torch.matmul(A, B))
# tensor([[19., 22.],
#         [43., 50.]])

# 等价写法
print(A @ B)
print(A.mm(B))  # 仅适用于2D

# 批量矩阵乘法
A_batch = torch.randn(10, 3, 4)
B_batch = torch.randn(10, 4, 5)
C_batch = torch.bmm(A_batch, B_batch)  # (10, 3, 5)

# 转置
print(A.T)
# tensor([[1., 3.],
#         [2., 4.]])

# 逆矩阵
print(torch.inverse(A))

# 行列式
print(torch.det(A))  # tensor(-2.0000)

聚合运算​

x = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.float)

# 求和
print(torch.sum(x))        # tensor(21.)
print(torch.sum(x, dim=0)) # tensor([5., 7., 9.]) 按列
print(torch.sum(x, dim=1)) # tensor([6., 15.]) 按行

# 均值
print(torch.mean(x))        # tensor(3.5000)
print(torch.mean(x, dim=1)) # tensor([2., 5.])

# 最大值、最小值
print(torch.max(x))         # tensor(6.)
print(torch.max(x, dim=1))  # 返回值和索引

# 标准差、方差
print(torch.std(x))
print(torch.var(x))

# 乘积
print(torch.prod(x))

# 范数
print(torch.norm(x))        # L2范数
print(torch.norm(x, p=1))   # L1范数

广播机制​

当两个张量形状不同时，PyTorch 会自动进行广播：

x = torch.tensor([[1, 2, 3], [4, 5, 6]])  # (2, 3)
y = torch.tensor([10, 20, 30])            # (3,)

# 广播后 y 变成 (2, 3)
print(x + y)
# tensor([[11, 22, 33],
#         [14, 25, 36]])

# 广播规则：
# 1. 从右向左比较维度
# 2. 维度相等或其中一个为1
# 3. 缺失的维度视为1

x = torch.randn(2, 1, 3)  # (2, 1, 3)
y = torch.randn(4, 3)     # (4, 3)
z = x + y                 # (2, 4, 3)

索引与切片​

基本索引​

x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

# 单个元素
print(x[0, 1])  # tensor(2)

# 行索引
print(x[0])     # tensor([1, 2, 3])
print(x[0, :])  # tensor([1, 2, 3])

# 列索引
print(x[:, 0])  # tensor([1, 4, 7])

# 切片
print(x[0:2, 1:3])
# tensor([[2, 3],
#         [5, 6]])

# 步长
print(x[::2, ::2])
# tensor([[1, 3],
#         [7, 9]])

高级索引​

x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

# 布尔索引
mask = x > 5
print(mask)
# tensor([[False, False, False],
#         [False, False, True],
#         [True, True, True]])

print(x[mask])  # tensor([6, 7, 8, 9])

# 条件筛选
print(x[x > 5])  # tensor([6, 7, 8, 9])

# 整数索引
rows = torch.tensor([0, 2])
cols = torch.tensor([1, 2])
print(x[rows, cols])  # tensor([2, 9])

# torch.index_select
indices = torch.tensor([0, 2])
print(torch.index_select(x, dim=0, index=indices))
# tensor([[1, 2, 3],
#         [7, 8, 9]])

修改张量​

x = torch.zeros(2, 3)

# 修改单个元素
x[0, 0] = 1
print(x)

# 修改一行
x[1, :] = torch.tensor([4, 5, 6])
print(x)

# 条件修改
x[x == 0] = -1
print(x)

# scatter_ 方法
x = torch.zeros(2, 3)
index = torch.tensor([[0, 1, 2], [2, 0, 1]])
src = torch.tensor([[1, 2, 3], [4, 5, 6]])
x.scatter_(1, index, src)
print(x)
# tensor([[1., 2., 3.],
#         [5., 6., 4.]])

形状操作​

改变形状​

x = torch.arange(12)

# reshape（不改变原张量）
y = x.reshape(3, 4)
print(y.shape)  # torch.Size([3, 4])

# view（共享内存）
z = x.view(2, 6)
print(z.shape)  # torch.Size([2, 6])

# 注意：view 要求张量在内存中连续
# 如果不连续，需要先调用 contiguous()
x = x.reshape(3, 4).t()  # 转置后不连续
# y = x.view(2, 6)  # 报错
y = x.contiguous().view(2, 6)  # 正确

# 展平
x = torch.randn(2, 3, 4)
y = x.flatten()        # (24,)
y = x.flatten(1)       # (2, 12) 从第1维开始展平

# 增加维度
x = torch.tensor([1, 2, 3])  # (3,)
y = x.unsqueeze(0)  # (1, 3)
z = x.unsqueeze(1)  # (3, 1)

# 减少维度
x = torch.randn(1, 3, 1)
y = x.squeeze()     # (3,)
z = x.squeeze(0)    # (3, 1)

拼接与分割​

# 拼接
x = torch.tensor([[1, 2], [3, 4]])
y = torch.tensor([[5, 6], [7, 8]])

# 按行拼接
z = torch.cat([x, y], dim=0)
print(z)
# tensor([[1, 2],
#         [3, 4],
#         [5, 6],
#         [7, 8]])

# 按列拼接
z = torch.cat([x, y], dim=1)
print(z)
# tensor([[1, 2, 5, 6],
#         [3, 4, 7, 8]])

# 增加新维度拼接
z = torch.stack([x, y], dim=0)
print(z.shape)  # torch.Size([2, 2, 2])

# 分割
x = torch.arange(10)
a, b, c = torch.split(x, [3, 3, 4])
print(a, b, c)
# tensor([0, 1, 2]) tensor([3, 4, 5]) tensor([6, 7, 8, 9])

# 均匀分割
x = torch.arange(12).reshape(3, 4)
chunks = torch.chunk(x, 2, dim=1)
print(chunks[0].shape, chunks[1].shape)
# torch.Size([3, 2]) torch.Size([3, 2])

排列与转置​

x = torch.arange(24).reshape(2, 3, 4)

# 转置（仅适用于2D）
x_2d = torch.arange(6).reshape(2, 3)
print(x_2d.t())
# tensor([[0, 3],
#         [1, 4],
#         [2, 5]])

# 交换维度
print(x.transpose(0, 1).shape)  # torch.Size([3, 2, 4])

# 置换维度
print(x.permute(2, 0, 1).shape)  # torch.Size([4, 2, 3])

GPU 加速​

张量设备管理​

# 检查 CUDA 是否可用
print(torch.cuda.is_available())

# 获取 GPU 数量
print(torch.cuda.device_count())

# 创建张量时指定设备
x_cpu = torch.tensor([1, 2, 3])
x_gpu = torch.tensor([1, 2, 3], device='cuda')

# 张量移动到 GPU
x_gpu = x_cpu.to('cuda')
x_gpu = x_cpu.cuda()
x_gpu = x_cpu.to(device='cuda:0')

# 张量移动到 CPU
x_cpu = x_gpu.to('cpu')
x_cpu = x_gpu.cpu()

# 查看张量所在设备
print(x_cpu.device)  # cpu
print(x_gpu.device)  # cuda:0

多 GPU 操作​

# 指定 GPU
x = torch.tensor([1, 2, 3], device='cuda:0')
y = torch.tensor([4, 5, 6], device='cuda:1')

# 注意：不同 GPU 上的张量不能直接运算
# z = x + y  # 报错

# 需要先移动到同一设备
y = y.to('cuda:0')
z = x + y

性能对比​

import time

# CPU 计算
x_cpu = torch.randn(5000, 5000)
start = time.time()
for _ in range(100):
    y_cpu = x_cpu @ x_cpu.T
print(f"CPU 时间: {time.time() - start:.3f}s")

# GPU 计算
if torch.cuda.is_available():
    x_gpu = x_cpu.cuda()
    torch.cuda.synchronize()  # 等待 GPU 完成
    start = time.time()
    for _ in range(100):
        y_gpu = x_gpu @ x_gpu.T
    torch.cuda.synchronize()
    print(f"GPU 时间: {time.time() - start:.3f}s")

原地操作​

以 _ 结尾的操作是原地操作，直接修改张量：

x = torch.tensor([1, 2, 3])

# 非原地操作
y = x.add(1)  # 返回新张量
print(x)  # tensor([1, 2, 3])

# 原地操作
x.add_(1)  # 直接修改 x
print(x)  # tensor([2, 3, 4])

# 常见原地操作
x.fill_(0)      # 填充
x.zero_()       # 置零
x.mul_(2)       # 乘法
x.div_(2)       # 除法
x.scatter_()    # 散射
注意

原地操作会丢失梯度信息，在自动求导时需谨慎使用。

小结​

本章我们学习了：

  1. 张量概念：多维数组的泛化形式
  2. 创建张量：从列表、工厂函数、随机数、NumPy 创建
  3. 张量属性：形状、维度、数据类型、设备
  4. 张量运算：算术运算、矩阵运算、聚合运算
  5. 索引切片：基本索引、高级索引、条件筛选
  6. 形状操作：reshape、view、拼接、分割
  7. GPU 加速：设备管理、性能优化

参考资源​