PyTorch保存模型断点以及加载断点继续训练

本文链接：https://blog.csdn.net/qq_53682472/article/details/149718714

问题：
我在训练程序时忘记写中间保存模型了，现在我的程序使用命令 nohup python -u rf_agent_gpu_cubic.py >> output_rf_cubic.out 2>&1 & 在后台运行，我想要保存现在的模型，然后中断程序，那我之后如何接着训练
我的源代码中total_timesteps=50000，我想要变小一些，大概20000.
rl_agent.train_ppo(total_timesteps=50000)
请问这时我该怎么做，我不想前面的timesteps白训练了

解决：

一、代码中以及设置代码：PyTorch保存模型断点以及加载断点继续训练

参考博文pytorch训练中断后，如何在之前的断点处继续训练_模型训练中途中断可以继续吗-CSDN博客

二、针对上面代码中未写断点并保存模型的代码

1. 立即保存当前模型并中断程序

由于您的代码没有自动保存中间模型的功能，我们需要手动操作。请按以下步骤：

(1) 修改源代码添加模型保存功能

在您的 rf_agent_gpu_cubic.py 文件中添加以下代码：

import signal
import sys

# 添加信号处理函数（放在训练开始前）
def save_model(signum, frame):
    print("接收到中断信号，正在保存模型...")
    rl_agent.save("interrupted_model")  # 替换为您的模型保存方法
    sys.exit(0)

signal.signal(signal.SIGINT, save_model)  # 捕获Ctrl+C
signal.signal(signal.SIGTERM, save_model)  # 捕获kill命令

(2) 保存当前模型并结束进程

查找进程ID:

bash

ps aux | grep rf_agent_gpu_cubic.py

输出类似：

text

user  12345  0.0  0.0  00000  pts/0  R+   Jul28  0:00 python -u rf_agent_gpu_cubic.py

记录PID（示例中为 12345）

发送中断信号:

bash
```
kill -SIGINT 12345  # 触发保存模型
```
等待约10秒后检查目录下是否生成 interrupted_model 文件