如何利用GPU加速多分类转成对二分类数据集的筛选？

核心瓶颈分析 #

你当前的主要耗时点在于filter_task中的Python列表推导循环，以及重复传递完整数据集带来的冗余开销。以下是针对性的优化方案：

方案一：CPU端快速优化（无需GPU） #

通过Numpy向量操作替代Python循环，结合预分类索引预处理，可将数据集生成时间从半小时压缩至数分钟。

1. 替换Python循环为Numpy向量操作 #

原列表推导是纯Python循环，效率极低，改用Numpy的布尔向量操作（C实现）大幅提速：

import numpy as np
from itertools import repeat
import multiprocessing

def filter_task(args):
    X, y, ij, uniq_classes = args
    cls_a, cls_b = uniq_classes[ij[0]], uniq_classes[ij[1]]
    # 用Numpy向量操作生成掩码，替代Python循环
    mask = (y == cls_a) | (y == cls_b)
    X_pair = X[mask]
    y_pair = np.where(y[mask] == cls_a, 0, 1)
    return (ij, (X_pair, y_pair))

def make_pairs_ds(self, X, y):
    ijs = [[i, j] for i in range(self.n_cls-1) for j in range(i+1, self.n_cls)]
    # 用repeat避免重复生成X/y的引用列表，减少冗余
    with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool:
        self.pair_ds = pool.map(filter_task, zip(
            repeat(X), 
            repeat(y), 
            ijs,
            repeat(self.unique_cls)
        ))

2. 预分类索引预处理（进阶优化） #

提前遍历一次标签，记录每个类别的样本索引，后续生成配对数据集时直接合并索引，无需重复遍历整个标签数组：

def filter_task_with_preproc(args):
    X, idx_a, idx_b, ij = args
    combined_idx = np.concatenate([idx_a, idx_b])
    X_pair = X[combined_idx]
    # 直接生成标签，避免再次判断
    y_pair = np.concatenate([np.zeros(len(idx_a)), np.ones(len(idx_b))])
    return (ij, (X_pair, y_pair))

def make_pairs_ds(self, X, y):
    # 预处理：记录每个类别的样本索引（仅执行一次）
    class_indices = {}
    for i, cls in enumerate(self.unique_cls):
        class_indices[i] = np.where(y == cls)[0]
    
    ijs = [[i, j] for i in range(self.n_cls-1) for j in range(i+1, self.n_cls)]
    tasks = [(X, class_indices[i], class_indices[j], [i,j]) for i,j in ijs]
    
    with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool:
        self.pair_ds = pool.map(filter_task_with_preproc, tasks)

方案二：CuPy GPU加速（超大数据量场景） #

若数据集规模极大，可借助CuPy利用GPU并行能力进一步提速，尤其配合cuML的GPU版SVC可实现端到端GPU加速。

1. CuPy版本数据集生成 #

import cupy as cp
from itertools import repeat
import multiprocessing

def filter_task_cp(args):
    X, y, ij, uniq_classes = args
    cls_a, cls_b = uniq_classes[ij[0]], uniq_classes[ij[1]]
    mask = (y == cls_a) | (y == cls_b)
    X_pair = X[mask]
    y_pair = cp.where(y[mask] == cls_a, 0, 1)
    # 若后续用CPU版SVC，需转Numpy数组；用cuML则无需转换
    return (ij, (X_pair.get(), y_pair.get()))

def make_pairs_ds(self, X, y):
    # 提前将数据转为CuPy数组（GPU内存）
    X_cp = cp.array(X)
    y_cp = cp.array(y)
    uniq_classes_cp = cp.array(self.unique_cls)
    
    ijs = [[i, j] for i in range(self.n_cls-1) for j in range(i+1, self.n_cls)]
    with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool:
        self.pair_ds = pool.map(filter_task_cp, zip(
            repeat(X_cp), 
            repeat(y_cp), 
            ijs,
            repeat(uniq_classes_cp)
        ))

注意事项 #

• CuPy数组在多进程中传递时共享GPU显存，无需额外拷贝，比Numpy多进程更高效
• 若使用cuML的SVC，可直接用CuPy数组训练，避免CPU/GPU数据转换开销

关于Numba的补充说明 #

你之前用Numba变慢，大概率是未启用nopython模式或未针对向量场景优化。尝试以下写法可能改善，但通常Numpy向量操作已足够高效：

from numba import jit

@jit(nopython=True)
def create_mask(y, cls_a, cls_b):
    mask = np.zeros(len(y), dtype=np.bool_)
    for i in range(len(y)):
        if y[i] == cls_a or y[i] == cls_b:
            mask[i] = True
    return mask

def filter_task_numba(args):
    X, y, ij, uniq_classes = args
    cls_a, cls_b = uniq_classes[ij[0]], uniq_classes[ij[1]]
    mask = create_mask(y, cls_a, cls_b)
    X_pair = X[mask]
    y_pair = np.where(y[mask] == cls_a, 0, 1)
    return (ij, (X_pair, y_pair))

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