PyTorch 张量的基本操作

import torch

# 创建一维张量（向量）
tensor1d = torch.tensor([1.0, 2.0, 3.0])  # 浮点型向量
tensor1d_int = torch.tensor([1, 2, 3])     # 整型向量

# 创建二维张量（矩阵）
tensor2d = torch.tensor([[1, 2, 3], [4, 5, 6]])  # 2x3矩阵
tensor2d_float = torch.tensor([[1., 2.], [3., 4.]])  # 2x2浮点矩阵

# 创建三维张量
tensor3d = torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])  # 2x2x2张量

# 创建全零张量（常用于初始化）
zeros = torch.zeros(3, 4)  # 3行4列的全零矩阵
zeros_3d = torch.zeros(2, 3, 4)  # 2x3x4的三维零张量

# 创建全一张量
ones = torch.ones(2, 3)    # 2行3列的全一矩阵
ones_like = torch.ones_like(zeros)  # 创建与zeros形状相同的全一张量

# 创建随机张量（常用于权重初始化）
rand_tensor = torch.rand(5, 5)  # 5x5的随机矩阵，值在[0,1)之间均匀分布
randn_tensor = torch.randn(3, 3)  # 3x3的正态分布随机矩阵，均值0，方差1

# 创建单位矩阵（常用于线性变换初始化）
eye_matrix = torch.eye(4)  # 4x4的单位矩阵
eye_with_offset = torch.eye(3, 4, k=1)  # 3x4的矩阵，对角线偏移1位

# 创建等差数列张量
arange_tensor = torch.arange(0, 10, 2)  # [0, 2, 4, 6, 8]
linspace_tensor = torch.linspace(0, 1, 5)  # [0.0, 0.25, 0.5, 0.75, 1.0]

tensor = torch.rand(3, 4)

# 数据类型 - 决定张量存储的数值类型
print(tensor.dtype)    # 如torch.float32、torch.int64等
# 常用类型：torch.float32, torch.float64, torch.int32, torch.int64

# 存储设备 - 决定张量存储在cpu还是gpu上
print(tensor.device)   # 如'cpu'或'cuda:0'
# 可以通过.to(device)方法在设备间移动张量

# 形状 - 描述张量的维度信息
print(tensor.shape)    # 如torch.size([3, 4])
print(tensor.size())   # 同shape的另一种访问方式
print(tensor.ndim)     # 张量的维度数，这里是2

# 梯度计算标志 - 决定张量是否需要计算梯度
print(tensor.requires_grad)  # 默认为false
# 当设置为true时，pytorch会跟踪对此张量的所有操作以计算梯度

# 其他属性
print(tensor.is_leaf)  # 是否为叶节点（用户直接创建的张量）
print(tensor.is_cuda)  # 是否存储在gpu上
print(tensor.grad)     # 存储的梯度值（如果有）

a = torch.tensor([1, 2, 3], dtype=torch.float32)
b = torch.tensor([4, 5, 6], dtype=torch.float32)

# 加法（逐元素相加）
c = a + b          # 直接运算符
c = torch.add(a, b)  # 函数形式
c = a.add(b)       # 方法形式

# 原地加法（不创建新张量）
a.add_(b)  # 等价于 a += b

# 乘法（逐元素相乘）
d = a * b          # [1*4, 2*5, 3*6] = [4, 10, 18]
d = torch.mul(a, b)
d = a.mul(b)

# 除法（逐元素相除）
e = a / b          # [1/4, 2/5, 3/6]
e = torch.div(a, b)

# 矩阵乘法
mat_a = torch.rand(3, 4)  # 3x4矩阵
mat_b = torch.rand(4, 5)  # 4x5矩阵
mat_prod = torch.matmul(mat_a, mat_b)  # 3x5结果矩阵
mat_prod = mat_a @ mat_b  # python 3.5+支持的运算符

# 广播机制示例
a = torch.rand(5, 1, 4)
b = torch.rand(  1, 3)
c = a + b  # 广播后形状为(5,3,4)

tensor = torch.arange(12).reshape(3, 4)
"""
tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11]])
"""

# 基本索引
print(tensor[0])       # 第一行: tensor([0, 1, 2, 3])
print(tensor[-1])      # 最后一行: tensor([ 8,  9, 10, 11])
print(tensor[:, 1])    # 第二列: tensor([1, 5, 9])
print(tensor[1:3, 2])  # 第二三行的第三列元素: tensor([6, 10])

# 高级索引
print(tensor[[0, 2], [1, 3]])  # 第0行第1列和第2行第3列: tensor([1, 11])
print(tensor[tensor > 5])      # 所有大于5的元素: tensor([6, 7, 8, 9, 10, 11])

# 修改子张量
tensor[1:, 2:] = 0  # 将第二行开始的右下2x2子矩阵置零

# 改变形状（不改变数据）
tensor = torch.rand(4, 6)
reshaped = tensor.reshape(3, 8)  # 总元素数必须保持不变
viewed = tensor.view(2, 12)      # 与reshape类似，但对内存连续性有要求

# 转置
transposed = tensor.t()  # 仅适用于2d张量
permuted = tensor.permute(1, 0)  # 通用转置方法

# 增加/减少维度
tensor = torch.rand(3, 4)
unsqueezed = tensor.unsqueeze(0)  # 在0维增加1维，变为1x3x4
squeezed = unsqueezed.squeeze()   # 去除所有长度为1的维度

# 连接张量
a = torch.rand(2, 3)
b = torch.rand(2, 3)
cat_0 = torch.cat([a, b], dim=0)  # 沿0维连接，结果为4x3
cat_1 = torch.cat([a, b], dim=1)  # 沿1维连接，结果为2x6

# 分割张量
chunks = torch.chunk(tensor, 3, dim=0)  # 沿0维分成3块
splits = torch.split(tensor, 2, dim=1)  # 沿1维每2列分一块

import numpy as np

# 张量转 numpy
tensor = torch.ones(5)
numpy_array = tensor.numpy()  # 共享内存，修改一个会影响另一个

# 注意：gpu上的张量需要先移动到cpu
gpu_tensor = torch.rand(5, device='cuda')
cpu_tensor = gpu_tensor.cpu()
numpy_array = cpu_tensor.numpy()

# numpy 转张量
np_array = np.array([1, 2, 3])
torch_tensor = torch.from_numpy(np_array)  # 同样共享内存

# 不共享内存的创建方式
torch_tensor_copy = torch.tensor(np_array)  # 创建新副本

# 检查gpu是否可用
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"using device: {device}")

# 创建张量并移动到gpu
tensor = torch.rand(10, 10).to(device)

# 直接在目标设备上创建张量
gpu_tensor = torch.rand(5, 5, device='cuda')

# 在gpu上执行运算
result = tensor * 2  # 如果tensor在gpu上，result也会在gpu上

# 将结果移回cpu
cpu_result = result.cpu()

# 多gpu示例（需要多gpu环境）
if torch.cuda.device_count() > 1:
    tensor = tensor.to('cuda:1')  # 显式指定第二个gpu

# 创建需要计算梯度的张量
x = torch.tensor(2.0, requires_grad=true)
w = torch.tensor(1.0, requires_grad=true)
b = torch.tensor(0.5, requires_grad=true)

# 构建计算图
y = w * x + b  # 线性函数
z = y**2       # 非线性变换

# 反向传播计算梯度
z.backward()  # 计算dz/dx, dz/dw, dz/db

# 查看梯度
print(x.grad)  # dz/dx = 2*(w*x+b)*w = 2*(1*2+0.5)*1 = 5.0
print(w.grad)  # dz/dw = 2*(w*x+b)*x = 2*(1*2+0.5)*2 = 10.0
print(b.grad)  # dz/db = 2*(w*x+b)*1 = 2*(1*2+0.5)*1 = 5.0

# 梯度清零（重要！在训练循环中需要）
x.grad.zero_()
w.grad.zero_()
b.grad.zero_()

# 禁用梯度跟踪（用于推理阶段）
with torch.no_grad():
    y = w * x + b  # 不会在计算图中记录

# 视图操作共享内存
a = torch.rand(4, 5)
b = a.view(20)  # b与a共享内存
c = a[1:3, 2:4]  # c与a共享内存

# clone()创建独立副本
d = a.clone()  # d与a不共享内存

# 创建稀疏张量
indices = torch.tensor([[0, 1, 2], [2, 0, 1]])  # 非零元素的坐标
values = torch.tensor([3, 4, 5], dtype=torch.float32)  # 非零元素的值
shape = torch.size([3, 3])  # 稀疏张量的形状

sparse_tensor = torch.sparse_coo_tensor(indices, values, shape)
print(sparse_tensor.to_dense())  # 转换为稠密张量

# 分布式训练中的张量分片
import torch.distributed as dist

if dist.is_initialized():
    # 将张量分片到不同进程
    tensor = torch.rand(10, 10)
    shard = tensor.chunk(dist.get_world_size(), dim=0)[dist.get_rank()]

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
x = x.to(device)  # 将张量移动到指定设备

x = torch.tensor(2.0, requires_grad=true)
y = x**2
y.backward()  # 自动计算梯度
print(x.grad)  # 输出梯度值 4.0