如何并行化GPU上运行的PyTorch线性求解函数

你的核心问题在于每个样本的非NaN位置不一致，导致每次求解的输入维度不同，无法直接用常规批量操作，原代码的双重循环完全没有发挥GPU的并行计算能力。下面是针对这个问题的优化方案：

核心思路：按掩码分组批量计算 #

将掩码完全相同的样本归为一组，同一组内的样本可以共享相同的掩码后设计矩阵A_masked，从而实现批量线性求解，大幅减少循环次数，充分利用GPU的并行资源。

优化后的代码 #

import torch
from tqdm import tqdm
from datetime import datetime
from collections import defaultdict
import numpy as np

# 设备设置
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
torch.manual_seed(42)

# 张量尺寸定义
B, C = 500, 500
M, N = 100, 50

# 生成并迁移张量到设备
input_tensor = torch.randn(N, B, C).to(device)
host_tensor = torch.randn(M, B, C).to(device)
A = torch.randn(N, M).to(device)

# 随机插入NaN模拟缺失数据
probability = 0.2
random_indices = torch.rand(input_tensor.shape, device=device) < probability
input_tensor[random_indices] = float('nan')

t_start = datetime.now()

# 1. 将张量扁平化，方便样本索引处理
input_flat = input_tensor.view(N, -1)  # 形状变为 (N, B*C)
host_flat = host_tensor.view(M, -1)

# 2. 生成每个样本的掩码，并转换为可哈希的分组键
masks = ~torch.isnan(input_flat)  # 形状 (N, B*C)
# 将布尔掩码转为uint8数组再转成元组，作为哈希键
mask_keys = [tuple(mask.cpu().numpy().astype(np.uint8)) for mask in masks.T]

# 3. 按掩码分组，记录每组的样本索引
groups = defaultdict(list)
for idx, key in enumerate(mask_keys):
    groups[key].append(idx)

# 4. 遍历分组批量求解
for mask_key, indices in tqdm(groups.items(), desc="Processing groups"):
    # 还原当前组的掩码张量
    mask = torch.tensor(mask_key, dtype=torch.bool, device=device)
    # 获取当前组的输入样本
    group_input = input_flat[:, indices]
    # 对输入和设计矩阵做掩码处理
    input_masked = group_input[mask]
    A_masked = A[mask]
    A_trans = A_masked.T

    # 批量计算线性系统
    lhs = A_trans @ A_masked
    rhs = A_trans @ input_masked
    # torch.linalg.solve支持右侧为矩阵，一次性求解整组样本
    solutions = torch.linalg.solve(lhs, rhs)

    # 将结果写回目标张量
    host_flat[:, indices] = solutions

t_stop = datetime.now()
print(f"Inversion took {(t_stop - t_start).total_seconds():.2f}s")

优化点说明 #

• 分组批量计算：把掩码相同的样本合并处理，循环次数从原有的B*C（250000次）大幅减少到掩码的唯一数量（通常远小于250000），彻底释放GPU并行能力。
• 扁平化张量：将三维张量转为二维，简化样本索引的分组和结果回写操作。
• 哈希掩码分组：通过将布尔掩码转换为可哈希的元组，快速完成样本分组，确保同一组样本使用相同的掩码后设计矩阵。

额外优化建议 #

1. 鲁棒性处理：如果A_trans@A_masked出现奇异矩阵导致求解失败，可以改用最小二乘求解，结果与原逻辑一致但更稳定：

   solutions = torch.linalg.lstsq(A_masked, input_masked).solution
   

2. 减少CPU-GPU数据传输：如果掩码数量极大，可以尝试在GPU上完成掩码的哈希转换（比如将布尔张量转为整数编码），避免频繁的CPU-GPU数据交互。

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