如何将含100k数据点与50M稀疏特征的自定义数据集加载到PyTorch？

针对你拥有的100k数据点、50M稀疏特征的自定义数据集，下面是具体的加载方案，包括Scipy CSC矩阵转PyTorch稀疏张量的方法：

一、Scipy CSC矩阵转PyTorch稀疏张量 #

PyTorch目前主要支持**COO格式（坐标格式）**的稀疏张量，所以我们需要先把CSC矩阵转换成COO格式，再创建PyTorch稀疏张量。操作步骤很简单：

1. 用Scipy的csc_matrix.tocoo()方法将CSC矩阵转为COO格式
2. 提取COO矩阵的行/列索引和数据值，转换为PyTorch张量
3. 用torch.sparse_coo_tensor构建稀疏张量，可按需指定设备（CPU/GPU）

示例代码：

import scipy.sparse as sp
import torch

# 模拟你的CSC稀疏矩阵（行=样本，列=特征）
csc_mat = sp.random(100000, 50000000, format='csc', density=1e-6)

# 转换为COO格式
coo_mat = csc_mat.tocoo()

# 提取索引和数据，转为PyTorch张量
indices = torch.tensor([coo_mat.row, coo_mat.col], dtype=torch.long)
values = torch.tensor(coo_mat.data, dtype=torch.float32)

# 创建稀疏张量
sparse_tensor = torch.sparse_coo_tensor(indices, values, size=coo_mat.shape)

# 如果GPU可用，移至GPU加速
if torch.cuda.is_available():
    sparse_tensor = sparse_tensor.cuda()

⚠️ 注意：PyTorch稀疏张量默认是不可变的，若需合并重复索引可调用sparse_tensor.coalesce()；绝对不要轻易转成稠密张量，50M特征的稠密张量会直接耗尽内存。

二、大规模数据集的高效加载策略 #

100k样本+50M特征的规模，全量加载到内存可能会有压力，推荐结合PyTorch的Dataset和DataLoader实现分批次、按需加载：

1. 自定义优化版SparseDataset #

提前将CSC转成COO格式并预处理索引范围，避免每次取样本时重复转换，大幅提升加载速度：

from torch.utils.data import Dataset, DataLoader
import numpy as np

class OptimizedSparseDataset(Dataset):
    def __init__(self, coo_row, coo_col, coo_data, labels, num_samples, num_features):
        self.coo_row = coo_row
        self.coo_col = coo_col
        self.coo_data = coo_data
        self.labels = torch.tensor(labels, dtype=torch.float32)
        self.num_features = num_features
        
        # 预处理：记录每个样本对应的COO数据索引范围
        _, sample_split_indices = np.unique(coo_row, return_index=True)
        self.sample_ranges = np.append(sample_split_indices, len(coo_row))

    def __len__(self):
        return len(self.sample_ranges) - 1

    def __getitem__(self, idx):
        # 获取当前样本对应的COO数据片段
        start_idx = self.sample_ranges[idx]
        end_idx = self.sample_ranges[idx+1]
        
        rows = self.coo_row[start_idx:end_idx]
        cols = self.coo_col[start_idx:end_idx]
        vals = self.coo_data[start_idx:end_idx]
        
        # 构建单样本稀疏张量
        indices = torch.tensor([rows, cols], dtype=torch.long)
        values = torch.tensor(vals, dtype=torch.float32)
        sample_tensor = torch.sparse_coo_tensor(indices, values, size=(1, self.num_features))
        
        return sample_tensor, self.labels[idx]

2. 配置高效的DataLoader #

合理设置参数，最大化加载效率：

# 假设你已经有标签数组labels
dataset = OptimizedSparseDataset(
    coo_row=coo_mat.row,
    coo_col=coo_mat.col,
    coo_data=coo_mat.data,
    labels=labels,
    num_samples=100000,
    num_features=50000000
)

dataloader = DataLoader(
    dataset,
    batch_size=64,  # 根据GPU显存调整，避免溢出
    shuffle=True,   # 训练时开启打乱
    num_workers=4,  # 多进程加载，根据CPU核心数设置（比如4/8）
    pin_memory=True # GPU训练时开启，加速内存到GPU的数据传输
)

三、额外优化技巧 #

• 预处理保存COO数据: 提前将COO格式的行/列索引、数据保存为numpy数组（比如用np.savez），后续加载时直接读取，省去重复转换CSC的时间
• 使用稀疏张量API: 模型中尽量用PyTorch的稀疏操作（比如torch.sparse.mm做稀疏矩阵乘法），避免稠密转换
• 混合精度训练: GPU训练时开启torch.cuda.amp.autocast()，节省显存同时不损失精度
• 稀疏优化器: 若模型参数是稀疏的，可使用PyTorch支持稀疏参数的优化器（如SGD、Adam均支持）

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