拉力测试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)
相关推荐
专注API从业者25 分钟前
Python/Java 代码示例:手把手教程调用 1688 API 获取商品详情实时数据
java·linux·数据库·python
java1234_小锋32 分钟前
[免费]基于Python的协同过滤电影推荐系统(Django+Vue+sqlite+爬虫)【论文+源码+SQL脚本】
python·django·电影推荐系统·协同过滤
看海天一色听风起雨落1 小时前
Python学习之装饰器
开发语言·python·学习
小憩-1 小时前
【机器学习】吴恩达机器学习笔记
人工智能·笔记·机器学习
却道天凉_好个秋2 小时前
深度学习(二):神经元与神经网络
人工智能·神经网络·计算机视觉·神经元
UQI-LIUWJ2 小时前
unsloth笔记:运行&微调 gemma
人工智能·笔记·深度学习
XiaoMu_0012 小时前
基于Python+Streamlit的旅游数据分析与预测系统:从数据可视化到机器学习预测的完整实现
python·信息可视化·旅游
THMAIL2 小时前
深度学习从入门到精通 - 生成对抗网络(GAN)实战:创造逼真图像的魔法艺术
人工智能·python·深度学习·神经网络·机器学习·生成对抗网络·cnn
却道天凉_好个秋2 小时前
计算机视觉(八):开运算和闭运算
人工智能·计算机视觉·开运算与闭运算
无风听海2 小时前
神经网络之深入理解偏置
人工智能·神经网络·机器学习·偏置