PyTorch中如何对两范围笛卡尔积应用向量化函数并使用GPU加速

完全不修改原函数的前提下，无法借助GPU算力高效生成全量配对数据集，核心限制有两点：

• 原函数内部的np.cos是NumPy的CPU侧运算，完全不支持GPU张量作为输入：哪怕提前把x、y的配对数据放到GPU显存，传入原函数时会直接触发类型错误，毕竟NumPy无法直接读取GPU内存中的数据。
• 原函数是标量级的计算逻辑，没有做向量化适配：要计算全量配对值只能通过嵌套循环逐次传入单个x、y标量，这种逐点调用的模式会完全抵消GPU的并行优势，实际运行效率甚至比纯CPU计算还慢。

硬要在不修改原函数的前提下跑通流程也可行：每次把GPU上的配对数据拷贝回CPU、喂给原函数计算、再把结果拷贝回GPU存储。但这种模式会引入大量不必要的CPU-GPU传输开销，GPU算力完全没有参与计算，和“借助GPU算力”的初始需求完全相悖，没有实际使用价值。

调整逻辑非常简单，不需要改动原函数的核心计算逻辑，只需要把NumPy的运算替换为PyTorch原生的张量运算即可。调整后的函数会自动跟随输入张量所在的设备（CPU/GPU）执行计算，同时支持广播向量化逻辑，可以直接处理全量配对的网格数据：

import torch

# GPU适配版函数，计算逻辑、输出精度和原NumPy版本完全一致
def f_gpu(x, y):
    return torch.cos(x) + y

测试场景完整运行流程 #

针对测试用例x=y=[0,1,2]的场景，实际执行步骤为：

1. 定义x、y的取值张量，指定存储在GPU上

   device = "cuda" if torch.cuda.is_available() else "cpu"
   x_vals = torch.tensor([0, 1, 2], device=device)
   y_vals = torch.tensor([0, 1, 2], device=device)
   

2. 生成所有x、y的配对网格（共3*3=9个样本）

   X, Y = torch.meshgrid(x_vals, y_vals, indexing="ij")
   

3. 直接将整个网格张量传入f_gpu，即可一次性算出所有配对的函数值，整个计算过程完全在GPU上执行

   dataset = f_gpu(X, Y)
   

   最终输出的dataset是3*3规格的张量，存储在GPU上，对应所有x、y配对的计算结果。

性能参考 #

小数据量下（比如本次测试的9个样本），两种实现的耗时差异几乎不可感知；但当x、y的取值范围扩大到各1万个点（即1亿个配对样本）时，逐点CPU计算的耗时通常达到小时级，而调整后的GPU向量化计算仅需数十毫秒即可完成。

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