PyTorch_GPU加速测试

最新推荐文章于 2025-11-13 09:30:00 发布

原创最新推荐文章于 2025-11-13 09:30:00 发布 · 1w 阅读

29 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #PyTorch #GPU加速

PyTorch 专栏收录该内容

16 篇文章

订阅专栏

本文通过示例代码展示了PyTorch库中CPU与GPU在进行大规模矩阵运算时的性能差异。首先，代码初始化了两个随机张量，并在CPU上执行矩阵乘法，记录下运行时间。随后，将同样的操作在GPU上执行，对比了GPU和CPU的运行效率，展示了GPU在处理大规模数据时的显著优势。

部署运行你感兴趣的模型镜像

示例代码：

import torch
import time

print(torch.__version__)        # 返回pytorch的版本
print(torch.cuda.is_available())        # 当CUDA可用时返回True

a = torch.randn(10000, 1000)    # 返回10000行1000列的张量矩阵
b = torch.randn(1000, 2000)     # 返回1000行2000列的张量矩阵

t0 = time.time()        # 记录时间
c = torch.matmul(a, b)      # 矩阵乘法运算
t1 = time.time()        # 记录时间
print(a.device, t1 - t0, c.norm(2))     # c.norm(2)表示矩阵c的二范数

device = torch.device('cuda')       # 用GPU来运行
a = a.to(device)
b = b.to(device)

# 初次调用GPU，需要数据传送，因此比较慢
t0 = time.time()
c = torch.matmul(a, b)
t2 = time.time()
print(a.device, t2 - t0, c.norm(2))

# 这才是GPU处理数据的真实运行时间，当数据量越大，GPU的优势越明显
t0 = time.time()
c = torch.matmul(a, b)
t2 = time.time()
print(a.device, t2 - t0, c.norm(2))

运行结果：

1.0.0
True
cpu 0.20143413543701172 tensor(3496775.5000)
cuda:0 0.28623294830322266 tensor(141487., device='cuda:0')
cuda:0 0.007987499237060547 tensor(141487., device='cuda:0')

您可能感兴趣的与本文相关的镜像