PyTorch训练函数内保存state_dict/checkpoint时的checkpoint初始化问题

一、仅初始化一次checkpoint的实现方式 #

完全可以实现，只需要把checkpoint = {}的初始化逻辑移到循环外层即可，示例代码如下：

def train_epoch(self):
    # 仅在循环开始前初始化一次checkpoint
    checkpoint = {}
    for epoch in tqdm.trange(self.epoch, self.max_epoch, desc='Train Epoch', ncols=100):
        self.epoch = epoch
        self.train()
        if epoch % 1 == 0:
            self.validate(checkpoint) 
        # 直接更新checkpoint内的最新检查点字段
        checkpoint['checkpoint_latest'] = {
            'epoch': self.epoch,
            'arch': self.model.__class__.__name__,
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optim.state_dict(),
            'model_save_criteria': self.model_save_criteria
        }
        torch.save(checkpoint, self.model_pth)

你提到的state_dict()由模型持有也不会影响该实现：PyTorch的model.state_dict()返回的是当前模型参数的快照，只要每轮都重新赋值model_state_dict为最新调用state_dict()的结果，就能保证存的是当前轮的最新参数。

二、每轮初始化checkpoint={}的合理性 #

这个写法完全合理，原因如下：

• 功能无问题：你每次都会覆盖写入同一个检查点文件，不管是每轮新建字典还是复用旧字典，最终保存的都是当前轮的最新状态，不会出现功能错误。
• 避免冗余字段：每次新建空字典可以自动清理上一轮可能遗留的临时字段（比如validate逻辑里写入的调试信息、临时计算指标），不会把无用内容写入检查点文件，也避免了忘记更新旧字段导致存到错误历史值的问题。
• 性能无损耗：创建空字典的开销极小，完全不会影响训练效率。

两种写法的适用场景 #

• 如果你只需要保存最新的检查点：两种写法都可以，每轮初始化checkpoint={}的写法更省心，不需要考虑字段遗留问题。
• 如果你需要同时保存多份状态（比如同时存最新检查点+历史最优检查点）：更适合全局初始化一次checkpoint，直接新增不同的key存储对应状态即可，不需要每轮都手动把最优状态复制到新字典里。

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