拉力测试cuda pytorch 把 4070显卡拉满

import torch
import time


def stress_test_gpu(matrix_size=16384, duration=300):
    """
    对GPU进行压力测试，通过持续的矩阵乘法来最大化GPU利用率

    参数:
        matrix_size: 矩阵维度大小，增大可提高计算复杂度
        duration: 测试持续时间（秒）
    """
    # 检查CUDA是否可用
    if not torch.cuda.is_available():
        print("CUDA不可用，无法进行GPU压力测试！")
        return

    # 获取GPU设备
    device = torch.device("cuda:0")
    print(f"使用GPU: {torch.cuda.get_device_name(device)}")

    # 创建大型矩阵并移到GPU
    print(f"正在创建 {matrix_size}x{matrix_size} 的矩阵...")
    a = torch.randn(matrix_size, matrix_size, device=device)
    b = torch.randn(matrix_size, matrix_size, device=device)

    # 预热GPU
    print("预热GPU中...")
    for _ in range(3):
        c = torch.matmul(a, b)
    torch.cuda.synchronize()

    # 开始压力测试
    print(f"开始GPU压力测试，持续 {duration} 秒...")
    start_time = time.time()
    iterations = 0

    try:
        while time.time() - start_time < duration:
            # 执行矩阵乘法（计算密集型操作）
            c = torch.matmul(a, b)

            # 同步设备以确保计算完成
            torch.cuda.synchronize()
            iterations += 1

            # 每10秒打印一次进度
            if iterations % 10 == 0:
                elapsed = time.time() - start_time
                print(f"已运行 {elapsed:.1f}s / {duration}s，迭代次数: {iterations}")

    except KeyboardInterrupt:
        print("测试被用户中断")

    finally:
        # 清理资源
        del a, b, c
        torch.cuda.empty_cache()

        # 计算性能指标
        elapsed = time.time() - start_time
        print(f"\n压力测试完成!")
        print(f"总运行时间: {elapsed:.2f} 秒")
        print(f"总迭代次数: {iterations}")
        print(f"平均每秒迭代: {iterations / elapsed:.2f}")


if __name__ == "__main__":
    # 可调整矩阵大小以平衡性能和显存使用
    # 16384x16384 矩阵约占用2GB显存 (4070有12GB显存)
    stress_test_gpu(matrix_size=16384, duration=300)