以下CUDA设备泛指NVIDIA显卡 或 启用ROCm的AMD显卡
import torch
import time
def cuda_benchmark(device_id, N=1000000):
# 指定要使用的显卡设备
torch.cuda.set_device(device_id)
# 创建输入数据
data = torch.ones(N).cuda()
# 启动CUDA操作,并记录执行时间
start_time = time.time()
for i in range(10000):
data += 1
torch.cuda.synchronize() # 等待CUDA操作执行完成
end_time = time.time()
# 将结果从GPU内存下载到主机内存
result = data.cpu().numpy()
# 打印Benchmark结果和执行时间
print(f"Benchmark结果:{result[:10]}")
print(f"执行时间:{end_time - start_time} 秒")
if __name__ == '__main__':
# 测试第一块显卡
device_id = 0
cuda_benchmark(device_id,10000000)
import torch
import torch.nn as nn
import time
class CUDABenchmarkModel(nn.Module):
def __init__(self):
super(CUDABenchmarkModel, self).__init__()
self.fc = nn.Linear(10, 10).cuda()
def forward(self, x):
return self.fc(x)
def cuda_benchmark(device_ids, N=10000000):
# 创建模型
model = CUDABenchmarkModel()
model = nn.DataParallel(model, device_ids=device_ids)
# 创建输入数据
data = torch.ones(N, 10).cuda()
# 启动CUDA操作,并记录执行时间
start_time = time.time()
for i in range(10000):
output = model(data)
torch.cuda.synchronize() # 等待CUDA操作执行完成
end_time = time.time()
# 打印执行时间
print(f"执行时间:{end_time - start_time} 秒")
if __name__ == '__main__':
# 同时测试3块显卡
device_ids = [0, 1, 2]
cuda_benchmark(device_ids=device_ids)
import torch
import torch.nn as nn
import torch.distributed as dist
import torch.multiprocessing as mp
import time
def cuda_benchmark(device_id, N=10000000):
# 指定要使用的显卡设备
torch.cuda.set_device(device_id)
print(f"该GPU的核心数量为:{torch.cuda.get_device_properties(device_id).multi_processor_count}")
# 创建输入数据
data = torch.ones(N).cuda()
# 启动CUDA操作,并记录执行时间
start_time = time.time()
for i in range(10000):
data += 1
torch.cuda.synchronize() # 等待CUDA操作执行完成
end_time = time.time()
# 将结果从GPU内存下载到主机内存
result = data.cpu().numpy()
# 打印Benchmark结果和执行时间
print(f"Benchmark结果:{result[:10]}")
print(f"执行时间:{end_time - start_time} 秒")
def main(num):
# 初始化多进程
mp.spawn(run, args=(num,), nprocs=num)
def run(rank,world_size):
"""每个进程的入口函数"""
# 初始化进程组
dist.init_process_group("nccl", init_method="tcp://127.0.0.1:23456", rank=rank, world_size=world_size)
# 指定设备ID
device_id = rank
# 在多个GPU上并行执行操作
model = cuda_benchmark(device_id)
if __name__ == '__main__':
# 同时启用3个进程(一个进程对应一块显卡)
device_numbers = 3
main(device_numbers)