|
|
本帖最后由 我輩樹である 于 2023-12-12 13:39 编辑
要计算GPU的性能，需要排除掉pcie switch和host device sync的过程，最佳的方式是直接在显存内生成矩阵，然后算矩阵乘法。
chat给的代码。
- import torch
- import time
- # 确保 CUDA 可用
- if not torch.cuda.is_available():
- raise SystemError("CUDA is not available. Tensor Cores require a CUDA-enabled GPU.")
- # 选择 CUDA 设备
- device = torch.device("cuda")
- # 显存中创建大型随机矩阵
- size = 4096 # 你可以调整这个大小
- # cuda core
- # a = torch.randn(size, size, device=device)
- # b = torch.randn(size, size, device=device)
- # tensor core
- a = torch.randn(size, size, device=device).half()
- b = torch.randn(size, size, device=device).half()
- # 热身 GPU
- for _ in range(10):
- c = torch.matmul(a, b)
- # 开始计时
- start = time.time()
- for _ in range(5000):
- # 执行矩阵乘法
- c = torch.matmul(a, b)
- # 结束计时
- elapsed_time = time.time() - start
- print(f"Time taken for matrix multiplication: {elapsed_time} seconds")
|
|
310112100042806
发表于 2023-12-11 20:31
