使用GPU加速大数据列表的条件掩码运算的方法（基于PyTorch框架）

准备工具

• Python
• Torch模块
• CUDA
• GPU

背景描述

当我们需要对一个较大规模列表（list）按照一定条件对其中的元素进行提取并进行特定运算时，可以尝试条件掩码的方法，即预先生成掩码算子，再将掩码算子作用到原列表后对列表内的数据进行提取。这个过程在CPU上是可行的，但当数据量非常大时（例如将神经网络模型参数保存到列表中，此时的列表长度通常是上万甚至上百万个的），会使得CPU承受很大的负担，且运算速度非常慢。为了提高能源与时间两个维度上的利用效率，我们可以寻找一种命令方式，使得该过程能够在GPU平台上执行。下面用示例代码举例。

示例伪代码

假如预设了一个较大规模的神经网络模型’model’，我们试着将其所有权重取出并保存到GPU上进行掩码运算，方法如下：

import torch
import torch.nn as nn

...

# 将权重数据提取到GPU平台上，实现超快运算，且避免使用循环语句
weights = [param.data.clone().to(device) for name, param in model.named_parameters() if 'weight' in name]

# 省略训练过程的代码
...

# 设定特殊条件算子，进行掩码提取处理 (以下的func1,func2和target_data是伪代码表示，实际情况下换成所用函数即可）
for param_data in weights:
    func1 = f(param_data)
    func2 = g(param_data)
    target_data = h(param_data)

	# 掩码处理，假设条件为提取出target_data中符合func1 - func2 > 0的数据
    mask = (func1 - func2) > 0:

	#提取符合条件的数据
    target_params = target_data[mask]

提取出target_params的过程完全在GPU上执行，因此大大加快了运算的速度。

实际尝试

本人在长度接近50万的原列表上进行双条件掩码操作，转移到GPU上后，相比于CPU，运算速度提升了9倍以上！

该图为GPU的利用率对比，右侧的波包表示使用GPU实现条件掩码运算的情况。