模型权重更新存储位置及Checkpoint工作机制全流程问询

反向传播后权重的存储位置 #

反向传播完成权重更新后，更新后的权重默认仍然留在GPU显存中，不会自动移回CPU内存。因为后续的前向传播、反向传播都需要复用这些权重，来回在GPU和CPU之间迁移会消耗大量带宽、拖慢训练速度，所以PyTorch、TensorFlow这类主流框架都会让权重尽可能持续驻留在GPU显存里。

权重何时会移回CPU内存？ #

只有两种场景会触发权重从GPU到CPU的迁移：

• 手动执行设备切换命令：比如调用model.to('cpu')（PyTorch）或tf.device('/CPU:0')（TensorFlow）主动将模型/权重转移到CPU。
• Checkpoint保存的临时操作：部分框架（比如PyTorch）在保存权重到磁盘时，会先把GPU上的权重张量临时拷贝到CPU内存，完成序列化写入磁盘后，CPU里的临时权重拷贝会被垃圾回收释放。不过TensorFlow等框架的部分保存方式可以直接从GPU显存读取权重写入磁盘，不需要经过CPU中转。

完整训练与Checkpoint执行流程 #

1. 初始化阶段：模型权重初始加载后，通过手动命令（如model.to('cuda')）或框架自动分配，将权重从CPU内存拷贝到GPU显存，后续所有计算都在GPU上进行。
2. 前向传播：GPU从显存中读取权重，执行对应层的kernel计算，输出预测结果并计算损失值。
3. 反向传播：基于损失值计算各权重的梯度，GPU直接在显存中用梯度更新权重参数——更新后的权重全程留在GPU显存，无自动回传CPU的操作。
4. Checkpoint保存触发：
   • 若使用PyTorch：调用torch.save(model.state_dict(), 'ckpt.pth')时，框架会先将GPU显存中的权重拷贝到CPU内存，序列化后写入磁盘文件；保存完成后，CPU中的临时权重数据会被自动清理。
   • 若使用TensorFlow：调用model.save()默认采用SavedModel格式时，可直接从GPU显存读取权重写入磁盘，无需中转CPU；若指定保存为HDF5格式（save_format='h5'），则会先将权重移到CPU再处理。
5. 后续训练循环：继续从GPU显存读取更新后的权重，重复前向-反向-更新流程，除非手动加载外部Checkpoint，否则不会从CPU或磁盘读取权重。

关于Checkpoint读取权重的补充 #

当你加载Checkpoint时（比如model.load_state_dict(torch.load('ckpt.pth'))），默认会把权重加载到CPU内存，之后需要手动调用model.to('cuda')再将权重移回GPU显存；也可以在加载时指定设备，比如torch.load('ckpt.pth', map_location='cuda')，直接把权重加载到GPU显存，跳过CPU中转。

内容的提问来源于stack exchange，提问作者user3696282