1. 从零上手PyTorch:为什么它成了深度学习项目的第一选择
如果你刚开始接触深度学习,或者从其他框架迁移过来,大概率会听到身边人反复提到一个名字——PyTorch。它不是一个新东西,但近几年在学术论文、工业项目和开源社区里的存在感越来越强。我最初从静态图框架转过来的时候,最大的感受是:写代码像在写普通的Python程序,而不是在“搭建一张图”。这种体验上的差异,直接决定了你调试模型时是花十分钟还是花两个小时。
PyTorch本质上是一个基于张量计算和自动求导的深度学习框架。张量你可以理解成多维数组,和NumPy的ndarray很像,但多了两个关键能力:一是可以在GPU上运行,二是能自动计算梯度。自动求导意味着你只需要写前向传播的逻辑,反向传播的链式法则由框架帮你完成。这对于搭建复杂模型来说,省掉的不仅是代码量,更是心智负担。
它适合谁?如果你在做研究、做原型验证、或者需要频繁修改模型结构,PyTorch的动态图机制会让你觉得很顺手。如果你是在做生产部署,PyTorch这几年也在补齐这块短板,TorchScript和ONNX导出让模型可以脱离Python环境运行。所以不管你是学生、算法工程师还是想入门AI的开发者,PyTorch都是一个值得投入时间的选择。
这篇文章不会只给你一堆API文档的复述,而是从实际使用的角度,把环境配置、张量操作、自动求导、模型搭建、训练循环、GPU加速、模型保存与加载这些环节里最容易踩坑的地方讲清楚。每个部分我都会解释“为什么这么做”,以及“不这么做会怎样”。你可以在本地跟着操作,也可以只看其中的某个章节解决当前遇到的问题。
2. 环境搭建:别让版本兼容问题成为你的第一个拦路虎
2.1 安装方式的选择逻辑
PyTorch的安装方式主要有两种:pip和conda。选哪个不是随便决定的。如果你用conda管理环境,conda install pytorch会同时帮你处理好CUDA运行时库的依赖,省去手动配置的麻烦。但conda的包更新速度有时会慢于pip,而且conda环境体积偏大。pip install torch的好处是轻量、更新快,但CUDA相关的库需要你自己确保系统里已经装好了对应版本的驱动和工具包。
我的建议是:如果你在个人电脑上做实验,用conda创建一个独立环境,把PyTorch和CUDA工具包一起装进去,这样不会污染系统环境。如果你在服务器上,而且服务器已经有全局的CUDA环境,用pip装PyTorch更干净。不管哪种方式,一定要去PyTorch官网的安装页面选择你的操作系统、包管理器、Python版本和CUDA版本,它会生成对应的安装命令。不要凭记忆手写命令,版本号差一个小版本就可能导致import失败。
注意:CUDA版本不是越新越好。你的显卡驱动版本决定了你最高能用的CUDA版本。用nvidia-smi命令查看驱动支持的CUDA版本,然后选择不超过这个版本的PyTorch CUDA构建。
2.2 验证安装是否真正可用
装完之后,很多人只运行import torch不报错就以为成功了。这不够。你需要验证三件事:PyTorch版本、CUDA是否可用、以及能否在GPU上实际执行一次运算。
import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else "No GPU") x = torch.randn(3, 3).cuda() print(x.device)如果torch.cuda.is_available()返回False,但你有GPU,常见原因有三个:一是装的是CPU版本的PyTorch,二是CUDA版本和驱动不匹配,三是环境变量里CUDA_VISIBLE_DEVICES被设置成了空。排查顺序就是从这三个方向逐一确认。我遇到过最隐蔽的一次是conda环境里装了一个CPU版本的torch,但pip list显示的名字一模一样,最后用torch.version.cuda查看编译时的CUDA版本才定位到问题。
2.3 虚拟环境的必要性
不管你用什么方式安装,强烈建议在虚拟环境里操作。原因很简单:不同项目依赖的PyTorch版本可能不同,有的项目需要1.x的API,有的用2.x的新特性。如果全部装在系统Python里,升级一个包可能让另一个项目跑不起来。conda create -n myenv python=3.10然后conda activate myenv,这两步花不了两分钟,但能帮你省掉未来可能出现的依赖冲突排查时间。
3. 张量操作:PyTorch的地基怎么打才稳
3.1 张量创建与NumPy的互操作
张量是PyTorch里最基本的数据结构。创建方式有很多种,torch.tensor()从Python列表创建,torch.zeros()和torch.ones()创建全零全一张量,torch.randn()创建标准正态分布随机张量。这些和NumPy的接口几乎一一对应,所以如果你有NumPy基础,上手会很快。
但有一个关键区别:torch.tensor()会复制数据,而torch.from_numpy()和numpy_array共享内存。这意味着如果你修改了从NumPy转换来的张量,原来的NumPy数组也会变。这个特性在数据预处理流水线里很有用,可以避免不必要的内存拷贝,但如果你不知情,可能会在调试时被“莫名其妙”的数据变化搞晕。
import numpy as np import torch arr = np.array([1, 2, 3]) t = torch.from_numpy(arr) t[0] = 100 print(arr) # array([100, 2, 3])反过来,t.numpy()也是共享内存的。如果你想要一个独立的副本,用t.clone().numpy()或者t.detach().cpu().numpy()。
3.2 形状操作中最容易混淆的几个方法
view、reshape、permute、transpose这几个方法经常让人搞混。view要求张量在内存里是连续的,否则会报错。reshape则更灵活,如果内存不连续它会自动复制一份再改变形状。所以在不确定张量是否连续的时候,用reshape更安全,但view在连续张量上效率更高。
permute和transpose都是交换维度,但transpose一次只能交换两个维度,permute可以一次性重新排列所有维度。比如一个形状为(B, C, H, W)的图像张量,你想变成(B, H, W, C),用permute(0, 2, 3, 1)一步到位。用transpose需要多次调用,而且容易搞错顺序。
还有一个常用的操作是squeeze和unsqueeze。squeeze去掉大小为1的维度,unsqueeze在指定位置插入一个大小为1的维度。在模型输入输出处理时特别常见,比如模型输出形状是(1, 10),但你需要(10,),用squeeze(0)就能去掉batch维度。
3.3 广播机制:省内存但也容易出错
PyTorch的广播机制和NumPy一致。当两个张量形状不同但满足广播规则时,框架会自动扩展维度进行运算。比如形状(3, 1)和(1, 4)相加,结果形状是(3, 4)。这很方便,但也是bug的温床。
我踩过的一个坑:两个张量本来应该逐元素相乘,结果因为形状分别是(N, 1)和(N,),广播之后变成了(N, N)的矩阵。这种错误不会报错,但计算结果完全不对。避免的方法是养成习惯,在关键运算前打印张量形状,或者用assert语句检查。
a = torch.randn(5, 1) b = torch.randn(5) c = a * b # 形状变成(5, 5),不是逐元素相乘如果你想要逐元素相乘,应该先把b变成(5, 1)或者把a变成(5,)。
4. 自动求导机制:理解它才能用好它
4.1 计算图是怎么构建的
PyTorch的自动求导核心是计算图。当你对一个设置了requires_grad=True的张量做运算时,PyTorch会记录下每一步操作,形成一个有向无环图。调用.backward()时,它从根节点反向遍历,用链式法则计算每个叶子节点的梯度。
这个机制的关键在于:计算图是在前向传播过程中动态构建的。每次迭代都会重新构建一张新图。这和静态图框架不同,静态图是先定义再运行。动态图的好处是你可以用普通的Python控制流,比如if和for,来改变模型结构,调试时也能逐行执行看中间结果。
x = torch.tensor([2.0], requires_grad=True) y = x ** 2 + 3 * x + 1 y.backward() print(x.grad) # 2*x + 3 = 74.2 梯度累积与清零的常见误区
在训练循环里,每次调用backward()之前必须把之前的梯度清零,否则梯度会累加。这就是为什么标准训练模板里有一行optimizer.zero_grad()。如果你忘了这一行,模型参数更新会变得混乱,loss可能不降反升。
但梯度累积也有它的用途。当显存不够大,无法设置较大的batch size时,你可以用多个小batch分别计算梯度但不立即更新参数,等累积到一定步数再更新。这样等效于大batch训练。实现方式就是在循环里控制zero_grad()和step()的调用频率。
accumulation_steps = 4 for i, (inputs, labels) in enumerate(dataloader): outputs = model(inputs) loss = criterion(outputs, labels) / accumulation_steps loss.backward() if (i + 1) % accumulation_steps == 0: optimizer.step() optimizer.zero_grad()注意loss要除以累积步数,否则梯度会放大对应倍数。
4.3 detach和no_grad的区别与使用场景
detach()和torch.no_grad()都用于切断梯度计算,但用法不同。detach()返回一个新的张量,从当前计算图中分离出来,但后续对这个新张量的操作仍然会被记录(如果它被设置为需要梯度的话)。no_grad()是一个上下文管理器,在这个上下文里所有运算都不会被记录到计算图中。
在推理阶段,用torch.no_grad()包裹整个前向传播,可以显著减少内存占用并加速计算。在训练阶段,如果你需要从某个中间变量取值但不希望影响梯度计算,用detach()。比如计算准确率时,需要把预测的logits转成类别标签,这个转换过程不需要梯度,就应该用detach()。
提示:在验证集上评估模型时,一定要用model.eval()和torch.no_grad()。前者切换BatchNorm和Dropout的行为,后者关闭梯度记录。两者缺一不可。
5. 模型搭建:从nn.Module开始构建你的网络
5.1 nn.Module的继承与forward方法
PyTorch里所有模型都继承自nn.Module。你需要做两件事:在__init__里定义网络层,在forward里定义数据流向。框架会自动帮你处理参数注册、梯度计算、GPU迁移等底层细节。
import torch.nn as nn class SimpleNet(nn.Module): def __init__(self, input_dim, hidden_dim, output_dim): super().__init__() self.fc1 = nn.Linear(input_dim, hidden_dim) self.relu = nn.ReLU() self.fc2 = nn.Linear(hidden_dim, output_dim) def forward(self, x): x = self.fc1(x) x = self.relu(x) x = self.fc2(x) return x这里有一个容易忽略的点:__init__里定义的层会被自动注册为模型的子模块,model.parameters()能遍历到它们。如果你把层放在一个普通的Python列表里,它们不会被注册,优化器就找不到这些参数。需要用nn.ModuleList或者nn.Sequential来包装。
5.2 常用层的选择与参数含义
nn.Linear是最基础的全连接层,参数是输入特征数和输出特征数。nn.Conv2d用于卷积,关键参数是in_channels、out_channels、kernel_size、stride和padding。padding的计算有一个常用公式:输出尺寸 = (输入尺寸 + 2*padding - kernel_size) / stride + 1。如果你想让输出尺寸和输入尺寸一致,kernel_size=3时padding=1,kernel_size=5时padding=2。
nn.BatchNorm2d在卷积层之后使用,对每个通道做归一化。它的参数是通道数。nn.Dropout在训练时随机置零一部分神经元,防止过拟合,参数是置零概率。注意Dropout在eval模式下不生效。
损失函数方面,分类任务常用nn.CrossEntropyLoss,它内部已经包含了softmax操作,所以你的模型输出应该是原始的logits,不要再加softmax。回归任务用nn.MSELoss。这些细节如果搞错,loss曲线会表现得非常奇怪。
5.3 参数初始化:默认值不一定适合你的任务
PyTorch的层有默认的初始化方式,比如Linear层用Kaiming均匀初始化。大多数情况下默认值够用,但如果你发现训练初期loss不下降或者梯度爆炸,可以考虑手动初始化。
def init_weights(m): if isinstance(m, nn.Linear): nn.init.xavier_uniform_(m.weight) nn.init.zeros_(m.bias) model.apply(init_weights)Xavier初始化适合tanh激活函数,Kaiming初始化适合ReLU。选择逻辑是:保持前向传播时每一层的方差一致,避免信号在深层网络中衰减或放大。
6. 训练循环:把各个组件串起来
6.1 标准训练循环的每个环节
一个完整的训练循环包含以下步骤:前向传播、计算损失、反向传播、参数更新、梯度清零。这五步的顺序不能乱。反向传播之前必须已经计算了损失,参数更新之前必须已经计算了梯度,下一次前向传播之前必须清零梯度。
model.train() for epoch in range(num_epochs): for inputs, labels in train_loader: inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step()把数据移到GPU上这一步,很多人会忘记对labels也做.to(device)。如果模型在GPU上而labels在CPU上,计算loss时会报设备不匹配的错误。
6.2 学习率调度:什么时候该降,降多少
学习率是最重要的超参数之一。太大导致loss震荡不收敛,太小导致收敛太慢。常用的策略有StepLR、CosineAnnealingLR和ReduceLROnPlateau。
StepLR每隔固定epoch数把学习率乘以一个衰减因子,比如每30个epoch乘以0.1。CosineAnnealingLR让学习率按余弦曲线从初始值降到最小值,适合训练周期固定的场景。ReduceLROnPlateau监控验证集loss,如果连续几个epoch不下降就降低学习率,适合不确定训练多久的情况。
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode='min', factor=0.5, patience=5 ) # 在验证之后调用 scheduler.step(val_loss)我个人的经验是:如果训练轮数明确,用CosineAnnealingLR;如果不确定,用ReduceLROnPlateau。StepLR的衰减点需要手动调,不太灵活。
6.3 验证集评估与模型选择
每个epoch训练完之后,需要在验证集上评估模型性能。这一步要用model.eval()切换模式,用torch.no_grad()关闭梯度。评估指标根据任务选择,分类用准确率、F1分数,回归用MSE、MAE。
保存模型时不要只保存最后一个epoch的结果,而是保存验证集上表现最好的那个。实现方式是维护一个best_score变量,每次验证后比较,如果当前分数更好就保存模型参数。
if val_acc > best_acc: best_acc = val_acc torch.save(model.state_dict(), 'best_model.pth')保存state_dict而不是整个模型,这样加载时更灵活,不会因为类定义变化而出问题。
7. GPU加速:把该放上去的都放上去
7.1 设备管理的正确姿势
PyTorch用.to(device)把张量和模型移到GPU上。device的定义通常是torch.device('cuda' if torch.cuda.is_available() else 'cpu')。模型移到GPU用model.to(device),数据移到GPU用tensor.to(device)。
有一个常见的性能陷阱:在训练循环里反复调用.to(device)会产生额外的开销。正确的做法是在DataLoader的collate_fn里或者循环外就把数据准备好。另外,从GPU取回数据用.cpu(),但要注意这会同步等待GPU完成所有计算,频繁调用会拖慢速度。
7.2 多GPU训练的基本概念
如果你有多张GPU,可以用nn.DataParallel或nn.parallel.DistributedDataParallel。DataParallel实现简单,一行代码就能用,但效率不如DistributedDataParallel,因为它在每个forward pass里复制模型和收集梯度。
if torch.cuda.device_count() > 1: model = nn.DataParallel(model) model.to(device)DistributedDataParallel需要启动多个进程,每个进程控制一张GPU,配置稍复杂但扩展性更好。对于大多数个人项目和小规模实验,DataParallel够用了。如果是在多机多卡环境,直接上DistributedDataParallel。
7.3 显存不够时的排查思路
显存不足(OOM)是训练时最常见的问题之一。排查顺序是:先看batch size是不是太大,尝试减半;再看模型参数量,如果模型本身太大,考虑用梯度检查点或者混合精度训练;最后看是否有不必要的中间变量被保留在计算图中。
混合精度训练用torch.cuda.amp,把部分计算转成float16,能显著减少显存占用并加速。用法是用autocast上下文管理器包裹前向传播,用GradScaler处理梯度缩放。
scaler = torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): outputs = model(inputs) loss = criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()注意:混合精度训练时,有些操作对float16敏感,比如softmax和loss计算。autocast会自动处理这些,但如果你手动写了自定义的运算,需要确认是否兼容。
8. 模型保存与加载:别让训练成果白费
8.1 state_dict与整个模型的区别
保存模型有两种方式:保存state_dict和保存整个模型。state_dict是一个字典,键是层的名字,值是对应的参数张量。保存整个模型用torch.save(model, path),加载时用torch.load(path)。后者看起来更方便,但有几个问题:它依赖模型类的定义,如果类结构变了加载会失败;它把模型结构和参数绑在一起,文件更大;跨设备加载时可能出问题。
推荐的做法是只保存state_dict,加载时先实例化模型类,再调用load_state_dict。
# 保存 torch.save(model.state_dict(), 'model.pth') # 加载 model = SimpleNet(784, 256, 10) model.load_state_dict(torch.load('model.pth')) model.to(device) model.eval()8.2 跨设备加载的坑
在GPU上训练保存的模型,在CPU上加载时,torch.load默认会把参数加载到原来的设备上。如果目标机器没有GPU,会报错。解决方法是加map_location参数。
state_dict = torch.load('model.pth', map_location='cpu')反过来,在CPU上保存的模型在GPU上加载,需要先加载再调用model.to(device)。不要在load_state_dict之前就把模型移到GPU,因为state_dict里的张量还在CPU上,直接加载会设备不匹配。
8.3 断点续训:保存优化器状态
如果你需要中断训练后继续,光保存模型参数不够,还需要保存优化器的状态和当前的epoch数。优化器状态里包含了动量、自适应学习率等历史信息,不保存的话续训时这些信息会丢失,可能导致loss曲线出现跳变。
checkpoint = { 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'loss': loss, } torch.save(checkpoint, 'checkpoint.pth')加载时对应地恢复每个组件的状态。这个习惯在训练大模型时尤其重要,因为一次训练可能跑好几天,中间遇到意外中断是常有的事。
9. 数据加载与预处理:别让IO成为瓶颈
9.1 Dataset和DataLoader的分工
PyTorch的数据加载由两个类配合完成:Dataset负责定义单条数据的读取和预处理逻辑,DataLoader负责批量组装、打乱顺序和多进程加载。你需要继承torch.utils.data.Dataset,实现__len__和__getitem__两个方法。
from torch.utils.data import Dataset, DataLoader class MyDataset(Dataset): def __init__(self, data, labels, transform=None): self.data = data self.labels = labels self.transform = transform def __len__(self): return len(self.data) def __getitem__(self, idx): x = self.data[idx] y = self.labels[idx] if self.transform: x = self.transform(x) return x, y__getitem__里只处理单条数据,不要在这里做批量操作。批量组装由DataLoader的collate_fn完成,默认的collate_fn会把单条数据堆叠成batch。
9.2 num_workers的设置与陷阱
DataLoader的num_workers参数控制用几个子进程加载数据。设置为0表示在主进程里加载,大于0会启动子进程。子进程加载能并行化IO和预处理,加快训练速度。但num_workers不是越大越好,一般设置为CPU核心数或者略小于核心数。
在Windows上使用num_workers>0时,必须把训练代码放在ifname== 'main':保护块里,否则会无限递归创建子进程。这是Windows的多进程机制决定的,Linux上没这个问题。我见过不少人在Windows上调试时被这个坑卡住,报错信息也不直观。
9.3 数据增强的常见操作
数据增强是防止过拟合的有效手段。对于图像数据,常用的增强包括随机裁剪、随机翻转、颜色抖动、旋转等。torchvision.transforms提供了这些操作的实现。
from torchvision import transforms train_transform = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness=0.2, contrast=0.2), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ])注意ToTensor()把PIL图像或NumPy数组转成张量,同时把像素值从[0, 255]缩放到[0, 1]。Normalize的均值和标准差是ImageNet数据集的统计值,如果你在自己的数据集上训练,最好重新计算这两个值。验证集和测试集不要做随机增强,只做确定性的预处理,比如resize和归一化。
10. 常见报错与排查经验
10.1 形状不匹配类错误
这类错误最常见,报错信息通常是“shape mismatch”或“size mismatch”。排查方法是打印每一步的张量形状,找到第一个不匹配的地方。全连接层要求输入是二维的(batch_size, features),如果你的输入是四维的图像张量,需要先展平。用x.view(x.size(0), -1)可以把除batch维度外的所有维度展平。
卷积层的通道数不匹配也很常见。上一层的输出通道数必须等于下一层的输入通道数。如果你改了某一层的out_channels,记得同步修改下一层的in_channels。
10.2 设备不匹配类错误
报错信息是“Expected all tensors to be on the same device”。原因是模型和数据不在同一个设备上。检查方法是打印model参数所在的设备和输入数据所在的设备。确保在训练循环里对inputs和labels都调用了.to(device)。
还有一种情况是损失函数里包含了需要梯度的张量和不需要梯度的张量,但它们在设备上不一致。比如你手动计算了一个指标,放在了CPU上,然后和GPU上的loss相加。这种错误比较隐蔽,需要仔细检查每一步运算的输入。
10.3 梯度相关的问题
梯度为None通常是因为张量的requires_grad没有设置为True,或者计算图被意外切断了。检查方法是打印叶子节点的grad属性。如果梯度爆炸,loss会变成nan,解决方法是加梯度裁剪。
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)梯度裁剪在RNN和Transformer的训练里几乎是标配。max_norm的选择需要根据具体任务调,一般从1.0开始试。
10.4 内存泄漏与显存增长
训练循环里如果不断累积不需要的张量,显存会持续增长直到OOM。常见的原因是列表里保存了每个step的loss但没有detach,导致计算图一直被引用。解决方法是保存loss时用loss.item()或者loss.detach().cpu()。
另一个原因是验证阶段没有用torch.no_grad(),导致验证集的计算图也被构建和保留。养成习惯:训练用model.train(),验证用model.eval()加torch.no_grad()。
11. 从实验到部署的过渡思路
11.1 TorchScript的基本用法
TorchScript可以把PyTorch模型转成不依赖Python的中间表示,方便在C++环境或移动端部署。转换方式有两种:trace和script。trace通过运行一次示例输入来记录计算图,适合没有控制流的模型。script通过解析Python源码来转换,支持if和for等控制流。
example_input = torch.randn(1, 784) traced_model = torch.jit.trace(model, example_input) traced_model.save('traced_model.pt')trace的局限是它只记录实际执行到的路径,如果模型里有依赖于输入的条件分支,trace会丢失其他分支。这种情况下用torch.jit.script。
11.2 ONNX导出与跨框架推理
ONNX是一种开放的模型交换格式,可以把PyTorch模型导出后在其他推理引擎上运行。导出用torch.onnx.export,需要提供示例输入和输入输出的名字。
torch.onnx.export( model, example_input, 'model.onnx', input_names=['input'], output_names=['output'], dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}} )dynamic_axes参数指定哪些维度是动态的,通常batch维度需要设为动态,这样导出的模型可以接受不同batch size的输入。导出后可以用onnxruntime加载并推理,速度通常比原生PyTorch快,尤其是在CPU上。
11.3 模型量化入门
量化是把模型的权重和激活从float32转成int8,减少模型体积和推理延迟。PyTorch提供了动态量化和静态量化两种方式。动态量化最简单,一行代码就能用,适合LSTM和Linear层较多的模型。
quantized_model = torch.quantization.quantize_dynamic( model, {nn.Linear}, dtype=torch.qint8 )静态量化需要校准数据,精度通常更好,但流程更复杂。量化会带来一定的精度损失,需要在速度和精度之间做权衡。对于大多数个人项目,动态量化已经能带来明显的速度提升。
12. 我踩过的几个印象深刻的坑
第一个坑是关于DataLoader的shuffle参数。训练集需要shuffle=True,验证集和测试集应该shuffle=False。我刚开始时忘了设置,导致验证集的顺序固定,某些batch的分布不均匀,评估指标波动很大。后来统一在验证时设shuffle=False,指标稳定了很多。
第二个坑是BatchNorm的momentum参数。默认值是0.1,但在小batch size下这个值可能偏大,导致running_mean和running_var更新太快,验证集表现不稳定。如果batch size小于16,可以考虑把momentum降到0.01或者用GroupNorm替代。
第三个坑是学习率预热。在Transformer类模型里,训练初期直接用大学习率会导致loss爆炸。需要在前几百个step里从很小的学习率线性增加到初始学习率。这个技巧在微调预训练模型时尤其重要,因为预训练权重已经在一个较好的位置,大学习率会破坏它们。
第四个坑是随机种子的设置。如果你需要复现实验结果,必须在训练开始前设置所有随机源的种子:torch.manual_seed、np.random.seed、random.seed,以及CUDA的种子。但即使这样,由于GPU并行计算的非确定性,结果也可能有微小差异。完全复现需要设置torch.use_deterministic_algorithms(True),但会牺牲一些性能。
这些经验在官方文档里不会写,只有实际跑过几个项目之后才会慢慢积累。PyTorch的生态很活跃,遇到问题去论坛搜一下通常能找到答案,但理解问题背后的原因比复制粘贴解决方案更重要。希望这些内容能帮你在PyTorch的路上少走一些弯路。