同一代码中实例化多个CrabNet模型的报错问题及解决咨询

问题描述 #

在催化项目中使用CrabNet工具验证不同数据集子集的训练效果时，同一代码中实例化两个CrabNet模型，调用第二个模型的fit方法触发以下错误：

File ".conda\lib\site-packages\torch\optim\optimizer.py", line 271, in wrapper for pre_hook in chain(_global_optimizer_pre_hooks.values(), self._optimizer_step_pre_hooks.values()): AttributeError: 'SWA' object has no attribute '_optimizer_step_pre_hooks'

复现代码：

"""Basic usage of CrabNet regression on elasticity dataset."""
from crabnet.utils.data import get_data
from crabnet.data.materials_data import elasticity, example_materials_property
from crabnet.crabnet_ import CrabNet

train_df, val_df = get_data(elasticity, "train.csv", dummy=True)
train_df_2, val_df_2 = get_data(example_materials_property, "train.csv", dummy=True)

cb = CrabNet(mat_prop="elasticity")
cb.fit(train_df)
val_pred, val_sigma = cb.predict(val_df, return_uncertainty=True)

cbn = CrabNet(mat_prop="example_materials_property")
cbn.fit(train_df_2)
val_pred_2, val_sigma_2 = cbn.predict(val_df_2, return_uncertainty=True)

解决方案 #

方法1：重置PyTorch全局优化器钩子 #

在实例化第二个CrabNet模型前，清空PyTorch的全局优化器前置钩子，避免跨模型的状态冲突：

import torch

# 训练第一个模型
cb = CrabNet(mat_prop="elasticity")
cb.fit(train_df)
val_pred, val_sigma = cb.predict(val_df, return_uncertainty=True)

# 重置全局优化器钩子
torch._C._clear_global_optimizer_pre_hooks()

# 训练第二个模型
cbn = CrabNet(mat_prop="example_materials_property")
cbn.fit(train_df_2)
val_pred_2, val_sigma_2 = cbn.predict(val_df_2, return_uncertainty=True)

方法2：用函数作用域隔离模型训练 #

将每个模型的训练逻辑封装到独立函数中，利用局部作用域隔离模型、优化器的状态，避免全局状态污染：

"""Basic usage of CrabNet regression on elasticity dataset."""
from crabnet.utils.data import get_data
from crabnet.data.materials_data import elasticity, example_materials_property
from crabnet.crabnet_ import CrabNet

def train_and_evaluate(mat_prop, train_data, val_data):
    model = CrabNet(mat_prop=mat_prop)
    model.fit(train_data)
    pred, sigma = model.predict(val_data, return_uncertainty=True)
    return model, pred, sigma

train_df, val_df = get_data(elasticity, "train.csv", dummy=True)
train_df_2, val_df_2 = get_data(example_materials_property, "train.csv", dummy=True)

# 训练第一个模型
cb, val_pred, val_sigma = train_and_evaluate("elasticity", train_df, val_df)

# 训练第二个模型
cbn, val_pred_2, val_sigma_2 = train_and_evaluate("example_materials_property", train_df_2, val_df_2)

方法3：升级PyTorch版本 #

该错误是PyTorch旧版本中SWA优化器与全局钩子机制的兼容性问题，升级PyTorch到1.13.0及以上版本后，该问题通常会被官方修复。

内容的提问来源于stack exchange，提问作者David Siret Marqués