Stylegan3如何保存指定tick状态到文件并从该文件恢复训练

一、保存指定tick的完整训练状态 #

StyleGAN3官方训练脚本内置了训练状态持久化逻辑，无需自行实现PyTorch权重保存逻辑：

• 训练启动参数中的--snap参数控制每多少个tick生成一次快照，例如设置--snap=5即每5个tick存一次状态。每次快照会生成两类文件：
  • network-snapshot-<tick编号>.pkl：仅存储推理所需的生成器、判别器权重，无法用于恢复训练
  • training-state-<tick编号>.pt：包含完整训练状态，覆盖当前tick数、迭代步数、优化器状态、随机数种子状态、损失函数累积状态等，是恢复训练所需的目标文件
• 若需要仅保存特定tick的状态，可以直接修改train.py训练循环内的快照触发逻辑，添加自定义判断即可，例如if cur_tick in [10, 25, 40]: 主动触发快照，无需修改其他底层逻辑。
• Google Colab环境训练时建议将输出目录--outdir设置为挂载后的谷歌云盘路径，避免虚拟机回收导致快照文件丢失。

二、从指定tick恢复训练 #

直接使用官方脚本内置的--resume参数，指定对应tick的training-state-*.pt文件路径即可：

python train.py --outdir=./training-runs --cfg=stylegan3-r --data=./你的数据集.zip \
  --gpus=1 --batch=32 --gamma=2 --snap=5 \
  --resume=/你的文件路径/training-state-00010.pt

• 注意不要使用.pkl文件恢复训练，该文件不包含优化器状态，恢复后优化器会重新初始化，直接导致训练曲线异常。
• 恢复训练后脚本会自动从指定tick的下一个进度开始运行，无需手动调整总训练轮次等参数，所有进度均会从状态文件中自动读取。

三、自定义保存/加载逻辑注意事项 #

若你需要自行实现状态保存逻辑，必须完整保存以下字段，否则会出现训练不连续的问题：

• 生成器G、判别器D的state_dict
• 生成器优化器G_opt、判别器优化器D_opt的state_dict
• 当前tick数cur_tick、当前累计训练图片数cur_nimg（StyleGAN3使用该字段计算训练进度）
• 所有随机数生成器状态：包括PyTorch的CPU、CUDA随机数状态，NumPy随机数状态
• 损失函数运行状态：例如路径长度正则的累积平均值等

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