import torch
import timedef 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)}")# 创建大型矩阵并移到GPUprint(f"正在创建 {matrix_size}x{matrix_size} 的矩阵...")a = torch.randn(matrix_size, matrix_size, device=device)b = torch.randn(matrix_size, matrix_size, device=device)# 预热GPUprint("预热GPU中...")for _ in range(3):c = torch.matmul(a, b)torch.cuda.synchronize()# 开始压力测试print(f"开始GPU压力测试，持续 {duration} 秒...")start_time = time.time()iterations = 0try: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_timeprint(f"已运行 {elapsed:.1f}s / {duration}s，迭代次数: {iterations}")except KeyboardInterrupt:print("测试被用户中断")finally:# 清理资源del a, b, ctorch.cuda.empty_cache()# 计算性能指标elapsed = time.time() - start_timeprint(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)