使用Python API会导致相同的行为。然而，可以像下面这样显式地调用wandb.init()和wandb.finish():

import logging
import wandb
from ludwig.contribs.wandb import WandbCallback
from ludwig.datasets import adult_census_income
from ludwig.api import LudwigModel

config = {
    "input_features": [
        {
            "name": "age",
            "type": "number",
        }
    ],
    "output_features": [
        {
            "name": "income",
            "type": "category",
        }
    ],
}

wandb.init(project="experiment") 
m = LudwigModel(config=config, callbacks=[WandbCallback()], logging_level=logging.INFO)
m.experiment(
    dataset=adult_census_income.load(),
    random_seed=1,
)
wandb.finish()

然后它就可以正常工作了。

你好，@MarselScheer,
感谢你在使用Ludwig上的wandb进行经验报告。我对wandb的集成不是很熟悉，但我注意到wandb.finish()在wandb回调的实现中缺失了。也许这就是为什么你在第二个设置中没有遇到OOM的原因。

你能尝试打补丁#3710,看看在你的第一次设置中是否观察到相同的OOM吗？

首先，我认为这解决了问题，因为前两次运行的结果符合预期(第一次额外分配了5GB内存，花了1或2分钟完成。第二次在几秒钟内完成)。之后它就不起作用了吗？!
尽管如此，它并没有解决问题，但改善了情况，因为现在如果我用CTRL+C中止执行，内存会被释放，我不需要退出容器。总的来说，我认为如果回调调用wandb.init(),也应该调用wandb.finish()。所以从这个Angular 来看，你可以合并PR,尽管它本身并没有解决问题。
然而，我也尝试过修改回调函数，导致问题的那一行似乎是
ludwig/ludwig/contribs/wandb.py
第47行：| | wandb.save(os.path.join(experiment_directory, "*")) |
注解掉这一行可以解决我进行的所有运行中的OOM问题。
我还调查了wandb的运行情况，特别是那些不知何故成功完成的前两个运行(使用了wandb.save())。唯一的区别似乎是这两个运行有这个artifact-folder(只包含一个文件)

指向一个parquet文件，内容如下：

不确定你期望wandb.save()上传什么信息，但对我来说，目前看起来这个信息已经被记录下来了，wandb.save()并不是必要的。然而，我认为这实际上应该保存除了模型本身之外的其他信息。