调用iter(dataloader)卡住且报__main__相关AttributeError的问题

环境与现象 #

• 本地环境：Python 3.8.11，执行 dataiter = iter(train_loader) 时程序卡住，随后抛出以下错误：

Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/Users/kelvin/opt/anaconda3/envs/torch-gpu/lib/python3.8/multiprocessing/spawn.py", line 116, in spawn_main
    exitcode = _main(fd, parent_sentinel)
  File "/Users/kelvin/opt/anaconda3/envs/torch-gpu/lib/python3.8/multiprocessing/spawn.py", line 126, in _main
    self = reduction.pickle.load(from_parent)
AttributeError: Can't get attribute 'HrDataset' on <module '__main__' (built-in)>

• 对比环境：Python 3.7.15 的 Google Colab 中运行完全正常。

数据集规模 #

Length of Dataset is 1470
Full: 1470
Train: 940
Valid: 236
Test: 294

相关代码 #

DataLoader 初始化代码 #

import multiprocessing as mp

bs = 32
# num_cpu = 2
num_cpu = mp.cpu_count()

train_loader = DataLoader(train, batch_size=bs, shuffle=True, num_workers=num_cpu, pin_memory=True)
valid_loader = DataLoader(valid, batch_size=bs, shuffle=False, num_workers=num_cpu, pin_memory=True)
test_loader = DataLoader(test, batch_size=bs, shuffle=False, num_workers=num_cpu, pin_memory=True)

HrDataset 类代码 #

class HrDataset(Dataset):

    def __init__(self, file_path):
        print('HrDataset is loading {}'.format(file_path))
        df = pd.read_csv(file_path)
        self.df = df
        self.df = self.preprocessing(df)
        print("Preprocessing is completed")
        print('Length of HrDataset is {}'.format(len(self.df)))

    def __getitem__(self, idx):
        X = np.array(self.df.iloc[idx, 1:]).astype(np.float32)
        y = self.df.iloc[idx, 0]

        return X, y

    def __len__(self):
        return len(df)

    def preprocessing(self, df):
        for col in df.columns:
            if df.dtypes[col] == 'object':
                df[col] = df[col].fillna('NA')
                df[col] = df[col].astype('category')
                if len(df[col].cat.categories) > 2:
                    df = pd.get_dummies(df, columns=[col])
                else:
                    df[col] = LabelEncoder().fit_transform(df[col])
            else:
                df[col] = df[col].fillna(0)
        return df

1. 修复 __len__ 方法的变量错误 #

HrDataset 的 __len__ 方法中，返回的是局部变量 df 的长度，而非实例属性 self.df，这会导致变量未定义错误，在多进程模式下间接引发卡住或序列化异常。修改为：

def __len__(self):
    return len(self.df)

2. 解决多进程序列化问题（AttributeError） #

Python 3.8 中 macOS 默认使用 spawn 多进程启动方式，要求自定义 Dataset 类能被子进程正确序列化，有两种解决方式：

方式一：将 Dataset 类独立到模块 #

把 HrDataset 类单独写到一个 .py 文件（如 dataset_utils.py），在主脚本中导入：

from dataset_utils import HrDataset

方式二：用 if __name__ == '__main__': 包裹主逻辑 #

将所有执行代码（Dataset 实例化、DataLoader 创建、迭代操作等）放入该判断块内，避免子进程重复执行初始化逻辑：

if __name__ == '__main__':
    # 实例化数据集
    train = HrDataset('train.csv')
    valid = HrDataset('valid.csv')
    test = HrDataset('test.csv')
    
    # 初始化 DataLoader
    bs = 32
    num_cpu = mp.cpu_count()
    train_loader = DataLoader(train, batch_size=bs, shuffle=True, num_workers=num_cpu, pin_memory=True)
    
    # 执行迭代操作
    dataiter = iter(train_loader)
    # 后续训练/测试代码...

3. 临时调试方案 #

如果上述方法暂未生效，可先将 num_workers 设置为 0（单进程模式），验证核心逻辑是否正常：

train_loader = DataLoader(train, batch_size=bs, shuffle=True, num_workers=0, pin_memory=True)

确认单进程运行无误后，再结合前两种方法恢复多进程配置。

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