中断场景下Skorch GridSearchCV的保存与续跑方法问询

核心思路 #

通过Skorch的Checkpoint回调保存单模型的训练状态（参数、优化器、训练历史），同时用joblib保存GridSearchCV的整体搜索进度，实现从任意中断点恢复训练：

1. 用Checkpoint自动保存每个参数组合、每个交叉验证折的模型状态
2. 用LoadInitState在训练前加载已保存的状态，续跑未完成的模型训练
3. 用joblib持久化GridSearchCV实例，跳过已完成的参数组合

完整代码示例 #

1. 导入依赖库

import numpy as np
import joblib
import torch
import torch.nn as nn
from sklearn.datasets import make_classification
from sklearn.model_selection import GridSearchCV
from skorch import NeuralNetClassifier
from skorch.callbacks import Checkpoint, LoadInitState

2. 定义神经网络模型

class MyNet(nn.Module):
    def __init__(self, hidden_units=32):
        super().__init__()
        self.fc1 = nn.Linear(20, hidden_units)
        self.fc2 = nn.Linear(hidden_units, 2)
    
    def forward(self, X):
        X = torch.relu(self.fc1(X))
        return self.fc2(X)

3. 准备训练数据

X, y = make_classification(2000, 20, n_informative=10, random_state=42)
X = X.astype(np.float32)
y = y.astype(np.int64)

4. 配置断点续跑回调

# 保存模型状态的Checkpoint回调
# 用dataset_idx和fold占位符区分不同参数组合、不同交叉验证折的模型
cp = Checkpoint(
    f_params='model_params_{{dataset_idx}}_{{fold}}.pt',
    f_optimizer='model_optimizer_{{dataset_idx}}_{{fold}}.pt',
    f_history='model_history_{{dataset_idx}}_{{fold}}.json',
    dirname='checkpoints',  # 保存目录，自动创建
)

# 加载已保存状态的LoadInitState回调
load_cp = LoadInitState(
    f_params='model_params_{{dataset_idx}}_{{fold}}.pt',
    f_optimizer='model_optimizer_{{dataset_idx}}_{{fold}}.pt',
    dirname='checkpoints',
)

5. 初始化Skorch模型并绑定回调

net = NeuralNetClassifier(
    MyNet,
    max_epochs=10,
    lr=0.01,
    callbacks=[cp, load_cp],  # 绑定两个回调
    verbose=1,
)

6. 实现GridSearchCV断点续跑

# 定义网格搜索参数
param_grid = {
    'hidden_units': [16, 32, 64],
    'lr': [0.001, 0.01, 0.1],
}

# 网格搜索进度保存路径
gs_progress_path = 'grid_search_progress.pkl'

# 尝试加载已有的搜索进度，否则创建新实例
try:
    gs = joblib.load(gs_progress_path)
    print("加载已保存的网格搜索进度，继续执行")
except FileNotFoundError:
    gs = GridSearchCV(
        net,
        param_grid,
        cv=3,
        scoring='accuracy',
        n_jobs=-1,
        verbose=2,
    )
    print("创建新的网格搜索实例")

# 启动/继续训练
gs.fit(X, y)

# 保存当前搜索状态（每次训练后自动覆盖）
joblib.dump(gs, gs_progress_path)

关键注意事项 #

• 目录与路径：确保checkpoints目录可读写，多进程训练时需使用共享存储（如本地磁盘）
• 代码一致性：续跑时模型结构、参数网格、训练配置必须与中断前完全一致，否则会加载失败
• 状态覆盖：Checkpoint默认在每个epoch结束时更新保存文件，确保恢复时拿到最新的训练状态
• GridSearchCV进度：joblib保存的实例包含已完成的参数组合结果，加载后会自动跳过已执行的任务

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