以下CUDA设备泛指NVIDIA显卡 或 启用ROCm的AMD显卡

• 测试环境：
  • Distributor ID: Ubuntu
  • Description: Ubuntu 22.04.3 LTS
  • Release: 22.04
  • Codename: jammy

1.首先，简单使用torch.ones测试CUDA设备

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)

2.使用自带的CUDABenchmarkModel测试CUDA设备

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)

3.使用nccl多进程的方式测试CUDA设备

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)