PyTorch GPU实现为何远慢于Numpy？求性能优化方案

你的GPU版本耗时极高，核心问题出在数据传输开销和小任务的GPU调度成本：

• CPU与GPU之间的数据传输：每次调用GPU_cost时，都要把numpy数组转成PyTorch tensor并拷贝到GPU，计算完成后又要把结果转回CPU的numpy数组。这种跨设备的数据拷贝非常耗时，尤其是当计算量本身很小时，传输开销会远远超过GPU计算的耗时。
• GPU的调度与启动开销：GPU擅长处理大规模并行计算任务，对于你这种小数据量（1000样本、20维度）的简单运算，GPU的启动、调度成本占比极高，反而不如CPU直接计算高效。
• 重复创建tensor的额外开销：每次调用函数都重新创建GPU tensor，没有复用设备内存，也会增加不必要的开销。

针对以上问题，可以通过以下几点大幅提升效率：

1. 避免频繁跨设备数据传输 #

把输入数据提前放到GPU上，让函数直接接收GPU tensor进行计算，仅在最终需要结果时再转回CPU（如果不需要numpy格式，甚至可以完全留在GPU上）。

2. 关闭梯度计算 #

因为你的场景不需要反向传播，使用torch.no_grad()上下文管理器可以避免PyTorch构建计算图，减少内存占用和额外开销。

3. 复用设备tensor #

提前分配好GPU上的tensor内存，避免每次调用都重新创建。

import numpy as np
import torch
import time

ndof = 2
x_dim = 20
n_samples = 1000

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def NP_cost(x):
    x_t_minus1 = np.concatenate((x[:, :ndof], x[:, :(x_dim - ndof)]), axis=1)
    x_t_minus2 = np.concatenate((x_t_minus1[:, :ndof], x_t_minus1[:, :(x_dim - ndof)]), axis=1)
    k_order = x + x_t_minus2 - 2 * x_t_minus1
    return np.sum(k_order ** 2, axis=1)

# 优化后的GPU版本函数，直接接收GPU tensor
def GPU_cost_optimized(x_tensor):
    with torch.no_grad():
        x_t_minus1 = torch.cat([x_tensor[:,:ndof], x_tensor[:, :(x_dim - ndof)]], dim=1)
        x_t_minus2 = torch.cat([x_t_minus1[:,:ndof], x_t_minus1[:,:(x_dim - ndof)]], dim=1)
        k_order = x_tensor + x_t_minus2 - 2*x_t_minus1
        return torch.sum(k_order**2, dim=1)

# 生成数据并提前转到GPU
x = np.random.rand(n_samples, x_dim)
x_gpu = torch.tensor(x, dtype=torch.float32).to(device)

# 测试原GPU版本
start_time = time.time()
print(GPU_cost(x)[0])
print(f"原GPU_cost执行时间: {time.time() - start_time:.4f} seconds")

# 测试优化后的GPU版本
start_time = time.time()
result = GPU_cost_optimized(x_gpu).cpu().numpy()
print(result[0])
print(f"优化后GPU_cost执行时间: {time.time() - start_time:.4f} seconds")

# 如果不需要numpy结果，直接在GPU上操作，耗时会更低
start_time = time.time()
result_gpu = GPU_cost_optimized(x_gpu)
print(result_gpu[0].item())
print(f"GPU上直接计算执行时间: {time.time() - start_time:.4f} seconds")

如果你的实际应用场景中数据量依然很小，其实完全没必要用GPU，CPU版本的效率已经足够。只有当计算量（比如样本数提升到10万以上，或者维度大幅增加）足够大时，GPU的并行优势才能体现出来，此时优化后的GPU版本会远超CPU版本的速度。

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