Ray Tune报错ImplicitFunc过大，求助定位大引用对象

核心问题根源 #

你遇到的ImplicitFunc体积超标问题，核心原因是**objective函数作为闭包隐式捕获了全局作用域的大对象**，尤其是未被正确隔离的tch_train数据集。即使你用ray.put()存储了数据引用，全局原对象仍然会被Ray打包进Actor，这就是421MiB的主要来源。此外，代码中还有几处逻辑错误和低效操作，进一步加剧了问题：

• 循环内重复调用ray.get()，导致数据重复加载
• optimizer绑定了未定义的全局model，既捕获额外变量又引发逻辑错误
• 测试阶段引用了未定义的net变量

修复步骤与代码调整 #

1. 隔离闭包捕获的全局变量 #

将objective函数内需要的数据仅通过Ray引用获取，避免捕获全局大对象；同时把所有重复的ray.get()移到函数开头，只执行一次。

2. 修正训练逻辑错误 #

为每个fold的模型单独创建optimizer，避免绑定全局无效模型；修复测试阶段的未定义变量问题。

修改后的完整代码示例 #

模型定义（保持不变）

INPUT_DIM = tch_train.features.shape[1] - 1  # 移除样本权重对应的输入特征
OUTPUT_DIM = tch_train.labels.shape[1]

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(INPUT_DIM, OUTPUT_DIM)

    def forward(self, input):
        output = F.softmax(F.relu(self.fc1(input)), dim=1)
        return output

主函数与objective调整

import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.model_selection import KFold
import ray
from ray import tune
from ray.tune.search.optuna import OptunaSearch

K_FOLDS = 5
loss_function = nn.CrossEntropyLoss(reduction='none')
kfold = KFold(n_splits=K_FOLDS, shuffle=True)

# 存储大对象到Ray对象存储
fold_indices = [(train_ids, test_ids) for train_ids, test_ids in kfold.split(tch_train)]
fold_indices_ref = ray.put(fold_indices)
tch_train_ref = ray.put(tch_train)

def objective(config):
    # 仅在函数开头获取一次数据，避免重复加载与闭包捕获全局对象
    fold_indices = ray.get(fold_indices_ref)
    tch_train = ray.get(tch_train_ref)
    
    # 为每个fold单独创建模型和对应的optimizer
    fold_models = []
    optimizers = []
    for fold in range(K_FOLDS):
        model = Net().to("cuda")
        fold_models.append(model)
        optimizers.append(torch.optim.SGD(
            model.parameters(), 
            lr=config["lr"], 
            momentum=config["momentum"]
        ))

    # 训练循环
    while True:
        fold_losses = [0.0 for _ in range(K_FOLDS)]  # 提前初始化损失列表

        for fold in range(K_FOLDS):
            train_ids, test_ids = fold_indices[fold]
            
            # 构建数据加载器
            train_subsampler = torch.utils.data.SubsetRandomSampler(train_ids)
            test_subsampler = torch.utils.data.SubsetRandomSampler(test_ids)
            trainloader = torch.utils.data.DataLoader(
                tch_train, 
                batch_size=config["batch_size"], 
                sampler=train_subsampler
            )
            testloader = torch.utils.data.DataLoader(
                tch_train, 
                batch_size=config["batch_size"], 
                sampler=test_subsampler
            )
            
            optimizer = optimizers[fold]
            
            # 训练步骤
            for i, data in enumerate(trainloader, 0):
                features, targets = data
                inputs = features[:,1:]
                sample_weights = features[:,0]
                
                optimizer.zero_grad()
                outputs = fold_models[fold](inputs)
                loss = loss_function(outputs, targets) * sample_weights
                loss.mean().backward()
                optimizer.step()

            # 测试步骤（修复未定义的net变量）
            with torch.no_grad():
                for i, data in enumerate(testloader, 0):
                    features, targets = data
                    inputs = features[:,1:]
                    sample_weights = features[:,0]
                    
                    outputs = fold_models[fold](inputs)
                    fold_losses[fold] += (loss_function(outputs, targets) * sample_weights).sum() 
            
        # 报告平均损失
        tune.report({"averaged_CEL": sum(fold_losses) / float(K_FOLDS)})

Tune配置（保持不变）

search_space = {"lr": tune.loguniform(1e-4, 1e-2), "momentum": tune.uniform(0.1, 0.9)}
algo = OptunaSearch() 

tuner = tune.Tuner(
    objective,
    tune_config=tune.TuneConfig(
        metric="averaged_CEL",
        mode="min",
        search_alg=algo,
    ),
    run_config=ray.train.RunConfig(
        stop={"training_iteration": 5},
    ),
    param_space=search_space,
)
results = tuner.fit()
print("Best config is:", results.get_best_result().config)

额外优化建议 #

如果数据集极大，可以考虑用Ray Data替代原生PyTorch DataLoader，它能更高效地在Ray集群中分发数据，无需手动管理对象引用，进一步降低Actor体积。

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