训练完成1个epoch后出现AttributeError: 'NoneType'无state_dict属性问题

问题场景 #

已按仓库要求安装所有依赖：

conda install pytorch==1.8.0 torchvision==0.9.0 cudatoolkit=10.2 -c pytorch
pip install opencv-python==4.4.0.46 termcolor==1.1.0 yacs==0.1.8

但完成一个epoch训练后，程序在保存checkpoint时报错。

相关代码 #

import os
import torch
import torch.distributed as dist
import numpy as np
import pdb
 
try:
    # noinspection PyUnresolvedReferences
    from apex import amp
except ImportError:
    amp = None

def save_checkpoint(config, epoch, model, max_accuracy, optimizer, lr_scheduler, logger):
    save_state = {'model': model.state_dict(),
                  'optimizer': optimizer.state_dict(),
                  'lr_scheduler': lr_scheduler.state_dict(),
                  'max_accuracy': max_accuracy,
                  'epoch': epoch,
                  'config': config}
    if config.AMP_OPT_LEVEL != "O0":
        save_state['amp'] = amp.state_dict()

    save_path = os.path.join(config.OUTPUT, f'ckpt_epoch_{epoch}.pth')
    logger.info(f"{save_path} saving......")
    torch.save(save_state, save_path)
    logger.info(f"{save_path} saved !!!")

训练日志 #

loss 6.5611 (6.7661)    grad_norm 2.8034 (nan)  mem 12138MB
[2023-06-08 11:25:02 deit_small_patch14_mask56_224_alpha1](main.py 488): INFO Train: [0/300][4970/5004] eta 0:00:37 lr 0.000051 time 1.0590 (1.1059)loss 6.6133 (6.7656)    grad_norm 2.9514 (nan)  mem 12138MB
[2023-06-08 11:25:13 deit_small_patch14_mask56_224_alpha1](main.py 488): INFO Train: [0/300][4980/5004] eta 0:00:26 lr 0.000051 time 1.1264 (1.1059)loss 6.3637 (6.7652)    grad_norm 3.3315 (nan)  mem 12138MB
[2023-06-08 11:25:24 deit_small_patch14_mask56_224_alpha1](main.py 488): INFO Train: [0/300][4990/5004] eta 0:00:15 lr 0.000051 time 1.0718 (1.1058)loss 6.5798 (6.7649)    grad_norm 2.6179 (nan)  mem 12138MB
[2023-06-08 11:25:35 deit_small_patch14_mask56_224_alpha1](main.py 488): INFO Train: [0/300][5000/5004] eta 0:00:04 lr 0.000051 time 1.0397 (1.1058)loss 6.7300 (6.7645)    grad_norm 3.6217 (nan)  mem 12138MB
[2023-06-08 11:25:38 deit_small_patch14_mask56_224_alpha1](main.py 497): INFO EPOCH 0 training takes 1:32:13

错误栈信息 #

Traceback (most recent call last):
  File "main.py", line 645, in <module>
    main(config)
  File "main.py", line 237, in main
    logger,
  File "/home2/coremax/Documents/MixPro/utils.py", line 193, in save_checkpoint
    save_state['amp'] = amp.state_dict()
AttributeError: 'NoneType' object has no attribute 'state_dict'

问题原因 #

代码尝试导入NVIDIA Apex库的amp（自动混合精度）模块失败，导致amp变量被赋值为None。但训练配置中的AMP_OPT_LEVEL不是"O0"（即启用了混合精度训练），触发了amp.state_dict()的调用，从而抛出AttributeError。

解决方案 #

方案1：安装Apex模块 #

执行以下命令安装适配PyTorch 1.8.0的Apex版本：

git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

如果编译时出现问题，可尝试去掉编译加速参数：

pip install -v --disable-pip-version-check --no-cache-dir ./

方案2：修改代码兼容未安装Apex的情况 #

修改save_checkpoint函数，增加对amp是否为None的判断，避免空对象调用：

def save_checkpoint(config, epoch, model, max_accuracy, optimizer, lr_scheduler, logger):
    save_state = {'model': model.state_dict(),
                  'optimizer': optimizer.state_dict(),
                  'lr_scheduler': lr_scheduler.state_dict(),
                  'max_accuracy': max_accuracy,
                  'epoch': epoch,
                  'config': config}
    # 增加amp是否为None的判断
    if config.AMP_OPT_LEVEL != "O0" and amp is not None:
        save_state['amp'] = amp.state_dict()

    save_path = os.path.join(config.OUTPUT, f'ckpt_epoch_{epoch}.pth')
    logger.info(f"{save_path} saving......")
    torch.save(save_state, save_path)
    logger.info(f"{save_path} saved !!!")

如果不需要混合精度训练，也可以直接将配置中的AMP_OPT_LEVEL设置为"O0"，这样就不会触发任何amp相关的代码逻辑。

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