3. PyTorch 基础及张量操作：GPU 张量创建、设备转移、基本运算

说实话，做基因组注释模型优化，第一步就是得把 PyTorch 张量玩熟。说白了，整个模型的数据流动、梯度计算、甚至显存管理，底层都是张量在撑。我刚开始用 PyTorch 的时候，总觉得张量跟 NumPy 数组差不多，直到有一次在 GPU 上跑模型，发现数据还在 CPU 里没搬过来，那叫一个慢……嗯，从那以后，我每次写完代码第一件事就是检查设备。

今天我们就来彻底搞懂三件事：

• GPU 张量怎么创建？
• 数据在 CPU 和 GPU 之间怎么搬？
• 基本的张量运算怎么在 GPU 上跑？

我会穿插一些我踩过的坑，希望能帮你省下 debug 的时间。

3.1 张量创建：CPU vs GPU

PyTorch 里最核心的数据结构就是 torch.Tensor。创建一个张量很简单，但创建在哪个设备上，差别巨大。

3.1.1 在 CPU 上创建

默认情况，torch.tensor 或者 torch.zeros 创建的都是 CPU 张量。

import torch
x = torch.randn(3, 4)
print(x.device)  # 输出：cpu

这没啥悬念。但注意，如果模型需要 GPU 计算，你就要手动搬过去。我以前犯过一个低级错误：在 CPU 上建了一堆序列数据，然后循环里每次都 .cuda() 一次，结果内存反复分配，速度反而更慢。所以最佳实践是：直接创建在 GPU 上。

3.1.2 在 GPU 上创建

有两种推荐方式：

1. 使用 device 参数 – 这是最干净的做法。
2. 先创建后转移 – 偶尔用于数据预处理流程。

# 方式一：直接指定 device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
x_gpu = torch.randn(3, 4, device=device)
print(x_gpu.device)  # 输出：cuda:0（如果有显卡）

# 方式二：先创建 CPU，再转移（不推荐，除非不得已）
x_cpu = torch.randn(3, 4)
x_gpu2 = x_cpu.to(device)

3.2 设备转移：搬数据的关键操作

基因组注释模型里，常有混合精度的需求，或者需要把数据从 CPU 预处理完再搬到 GPU。设备转移的两个主力函数：.to() 和 .cuda()。

3.2.1 .cuda() vs .to()

我个人推荐 .to()，因为它更灵活。比如我经常需要把数据转成 float16 再送进模型：

x = torch.randn(3, 4)
x_half = x.to(device='cuda', dtype=torch.float16)
print(x_half.dtype, x_half.device)  # torch.float16, cuda:0

3.2.2 注意：.to() 不是原地操作！

这是个经典坑。我曾经写过这样的代码：

tensor = torch.randn(3, 4)
tensor.to('cuda')        # 返回新张量，但没赋值回去
print(tensor.device)     # 还是 cpu — 是不是很尴尬？

3.2.3 多 GPU 场景下的设备转移

如果你要使用 DataParallel 或 DistributedDataParallel，别把数据全塞到 GPU 0。建议每个进程只创建自己那部分数据：

local_rank = torch.distributed.get_rank()
device = torch.device(f'cuda:{local_rank}')
x = torch.randn(1000, 128, device=device)

这样做的好处是避免主卡显存溢出。我在一次 4 卡训练基因序列模型时，就因为没按 rank 分配数据，结果 GPU0 直接被干到 23GB，其他卡还空着……

3.3 基本运算：GPU 加速的威力

张量运算包括加减乘除、矩阵乘法、广播、聚合等。在 GPU 上，这些运算会被 CUDA 核心并行执行，速度比 CPU 快几十倍甚至上百倍。但要注意，运算的粒度决定了加速效果。

3.3.1 常见运算一览

3.3.2 一个 GPU 加速的简单例子

假设你要做一万次矩阵乘法：

import torch
import time

n = 2048
a_cpu = torch.randn(n, n)
b_cpu = torch.randn(n, n)

device = torch.device('cuda')
a_gpu = a_cpu.to(device)
b_gpu = b_cpu.to(device)

# CPU 运算
start = time.time()
for _ in range(100):
    c = torch.mm(a_cpu, b_cpu)
print('CPU 耗时:', time.time() - start)

# GPU 运算（第一次包含编译预热，一般不计时）
for _ in range(5):
    c = torch.mm(a_gpu, b_gpu)
torch.cuda.synchronize()
start = time.time()
for _ in range(100):
    c = torch.mm(a_gpu, b_gpu)
torch.cuda.synchronize()
print('GPU 耗时:', time.time() - start)

在我的 2080Ti 上，CPU 要 20 多秒，GPU 只要 0.4 秒。你看，差了 50 倍不止。

3.3.3 运算中的注意点：内存碎片与 dtype

基因组序列数据经常是 long/int 类型的索引，但 GPU 上矩阵乘法通常用 float32 或 float16。我遇到过一次模型 loss 爆炸，查了半天发现是 embedding 的 dtype 不一致，导致自动广播时偷偷转成了低精度。解决办法是统一用 .to(dtype=torch.float32)。

另外，频繁的小张量创建会导致 CUDA 内存碎片。我在处理可变长序列时，曾经反复创建长度不同的 mask 张量，结果显存一点点变多，最后 OOM。后来改用 torch.zeros(batch_size, max_len, device=device) 一次性分配，然后把对应位置置 1，碎片就没了。

3.4 本章知识体系

下面这张图概括了 GPU 张量操作的完整流程。我建议你一边看代码一边对照这张图，理解每步的作用。

3.5 避坑总结

我把碰到过的典型问题列出来，你写代码时多留个心眼：

• 忘记赋值：.to() 返回新张量，必须赋值回原变量。
• 混合设备运算：CPU 和 GPU 张量不能直接相加，会报错 RuntimeError: Expected all tensors to be on the same device。
• 频繁小张量：每次循环都 torch.tensor([i]).cuda() 会导致显存碎片，建议预分配。
• 忘记 synchronize：测量 GPU 性能时务必加 torch.cuda.synchronize()，否则时间不准。
• dtype 不一致：运算前检查 .dtype，尤其是混合精度场景。

def to_device(data, device):
    if isinstance(data, torch.Tensor):
        return data.to(device, dtype=torch.float32, non_blocking=True)
    elif isinstance(data, dict):
        return {k: to_device(v, device) for k, v in data.items()}
    # ... 其他类型处理
    return data

掌握这些基础操作，后面讲混合精度训练和显存优化的时候，你就能直接上手了。记住：设备管理是 GPU 加速的第一道门槛，跨过去，后面的路就顺了。