RTX 50 系(如 5070/5080/5090) 要跑 PyTorch 然后 CUDA ≥ 12.8,配套的 PyTorch 通常需要 2.6+(更稳妥的是 2.8+ 或 nightly/cu128 预发布),Python 推荐 3.9–3.11(3.10 是稳妥选择)。显卡驱动和 CUDA 工具链、PyTorch wheel 三者必须“版本匹配”。50系的显卡好多老版本用不了。

我的电脑是微星泰坦16AI RTX5080,最后的版本:

测试设备: NVIDIA GeForce RTX 5080 Laptop GPU
显存总量: 15.9 GB
CUDA 版本: 12.8
PyTorch 版本: 2.10.0.dev20251119+cu128

#这是我最后的版本,你也可以用命令查一下在conda的prompt里
(base) C:\Users\huruiyang>conda activate hu_3_10  #进环境,没创建就没有

(hu_3_10) C:\Users\huruiyang>python -c "import torch; print(torch.__version__)"
2.10.0.dev20251119+cu128

(hu_3_10) C:\Users\huruiyang>python -c "import torchvision; print(torchvision.__version__)"
0.25.0.dev20251118+cu128

(hu_3_10) C:\Users\huruiyang>python -c "import torchaudio; print(torchaudio.__version__)"
2.10.0.dev20251117+cu128

以下是针对RTX 50系显卡运行PyTorch的配置指南,内容已按技术逻辑重组并去除步骤化表述:

环境预检

执行nvidia-smi记录当前驱动版本和CUDA支持情况。

驱动与工具链更新

从NVIDIA官网下载最新显卡驱动(需明确支持CUDA 12.8+),安装后重启系统。再次验证nvidia-smi输出的驱动版本与CUDA兼容性。

环境管理策略

使用Conda或Mamba创建独立虚拟环境(示例Python 3.10)。避免混合使用pip和conda安装同一环境的包,PyTorch CUDA wheel建议通过pip指定官方源安装。

旧环境清理(没有搞错过就不管,看下面安装方案就行)

在目标环境中执行以下命令序列:

#创建深度学习环境,hu_3_10是环境名字
conda create -n hu_3_10 python=3.10

# 激活目标环境
conda activate hu_3_10  

# 卸载 PyTorch 系列(pip)(可选)   
# 如果你之前装过,装错了
pip uninstall -y torch torchvision torchaudio

# 清理 pip 缓存(可选)
pip cache purge

# 若用 conda 安装过 pytorch-cuda 等,用 conda 卸载 就是把你整个环境删了重新建立
conda remove --name hu_3_10 --all   # 如果你愿意直接删整个环境并重建(最干净)
# 或者删除单个包:
conda remove pytorch pytorch-cuda -y
# 清理 conda 缓存
conda clean --all -y

手动检查并删除残留目录(如site-packages/torch*)。清理后建议重启系统。

PyTorch安装方案

#prompt里面装

对于RTX 50系显卡,优先选择支持CUDA 12.8的预构建版本:

# 新建虚拟环境 名字叫torch50的环境
conda create -n torch50 python=3.10
conda activate torch50

# (可选)安装 CUDA toolkit 系统级(若你需要 nvcc),否则 PyTorch wheel 自带 runtimes也可跑:
# 从 NVIDIA 官网下载 CUDA 12.8 installer(按需),然后重启。

# 安装 PyTorch(Nightly cu128,适配 RTX50)
pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
#或者用清华源安装,我是用的这个
pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128 --extra-index-url https://pypi.tuna.tsinghua.edu.cn/simple

# 如果想用官方 stable + cu128(当 Stable 支持时),用官方给出的 index-url 替换。

当官方发布稳定版cu128支持时,替换为对应的stable版本index-url。

验证安装(我用的别人那找的)

通过Python交互环境执行:

import torch
import time
import numpy as np


def comprehensive_gpu_test():
    """ GPU 综合性能测试"""
    print("=== GPU 性能测试 ===\n")

    device = torch.device("cuda")
    print(f"测试设备: {torch.cuda.get_device_name(0)}")
    print(f"显存总量: {torch.cuda.get_device_properties(0).total_memory / 1024 ** 3:.1f} GB")
    print(f"CUDA 版本: {torch.version.cuda}")
    print(f"PyTorch 版本: {torch.__version__}\n")

    # 1. 基础矩阵运算性能
    print("1. 矩阵运算性能测试:")
    sizes = [1024, 2048, 4096, 8192]

    for size in sizes:
        # GPU 测试
        a = torch.randn(size, size, device=device)
        b = torch.randn(size, size, device=device)

        # 预热
        for _ in range(3):
            _ = torch.mm(a, b)
        torch.cuda.synchronize()

        # 正式测试
        start_time = time.time()
        for _ in range(10):
            result = torch.mm(a, b)
        torch.cuda.synchronize()
        end_time = time.time()

        avg_time = (end_time - start_time) / 10 * 1000  # ms
        flops = 2 * size ** 3 * 10  # 10次操作的总FLOPS
        tflops = flops / (end_time - start_time) / 1e12

        print(f"   {size}x{size}: {avg_time:.2f} ms/次, {tflops:.2f} TFLOPS")

    # 2. 混合精度测试
    print("\n2. 混合精度 (FP16) 测试:")
    size = 4096
    a_fp32 = torch.randn(size, size, device=device)
    b_fp32 = torch.randn(size, size, device=device)
    a_fp16 = a_fp32.half()
    b_fp16 = b_fp32.half()

    # FP32
    torch.cuda.synchronize()
    start = time.time()
    result_fp32 = torch.mm(a_fp32, b_fp32)
    torch.cuda.synchronize()
    fp32_time = time.time() - start

    # FP16
    torch.cuda.synchronize()
    start = time.time()
    result_fp16 = torch.mm(a_fp16, b_fp16)
    torch.cuda.synchronize()
    fp16_time = time.time() - start

    speedup = fp32_time / fp16_time
    print(f"   FP32: {fp32_time * 1000:.2f} ms")
    print(f"   FP16: {fp16_time * 1000:.2f} ms")
    print(f"   FP16 加速比: {speedup:.2f}x")

    # 3. 深度学习模拟测试
    print("\n3. 深度学习模拟测试:")

    # 模拟卷积网络
    conv_layers = [
        torch.nn.Conv2d(3, 64, 3, padding=1),
        torch.nn.Conv2d(64, 128, 3, padding=1),
        torch.nn.Conv2d(128, 256, 3, padding=1),
        torch.nn.Conv2d(256, 512, 3, padding=1),
    ]

    for layer in conv_layers:
        layer = layer.to(device)

    batch_sizes = [1, 8, 16, 32]
    input_size = 224

    for batch_size in batch_sizes:
        x = torch.randn(batch_size, 3, input_size, input_size, device=device)

        torch.cuda.synchronize()
        start = time.time()

        # 前向传播
        for layer in conv_layers:
            x = layer(x)
            x = torch.relu(x)
            x = torch.max_pool2d(x, 2)

        torch.cuda.synchronize()
        end = time.time()

        throughput = batch_size / (end - start)
        print(f"   Batch size {batch_size:2d}: {(end - start) * 1000:.2f} ms, {throughput:.1f} imgs/sec")

    # 4. 内存带宽测试
    print("\n4. 显存带宽测试:")
    sizes_mb = [100, 500, 1000, 2000]

    for size_mb in sizes_mb:
        elements = size_mb * 1024 * 1024 // 4  # float32
        data = torch.randn(elements, device=device)

        torch.cuda.synchronize()
        start = time.time()

        # 内存复制测试
        for _ in range(10):
            data_copy = data.clone()

        torch.cuda.synchronize()
        end = time.time()

        bandwidth = (size_mb * 10 * 2) / (end - start) / 1024  # GB/s (读+写)
        print(f"   {size_mb} MB: {bandwidth:.1f} GB/s")

    # 5. Transformer 模拟测试
    print("\n5. Transformer 注意力机制测试:")

    seq_lengths = [512, 1024, 2048]
    d_model = 768
    num_heads = 12

    for seq_len in seq_lengths:
        # 模拟多头注意力
        batch_size = 8
        q = torch.randn(batch_size, num_heads, seq_len, d_model // num_heads, device=device)
        k = torch.randn(batch_size, num_heads, seq_len, d_model // num_heads, device=device)
        v = torch.randn(batch_size, num_heads, seq_len, d_model // num_heads, device=device)

        torch.cuda.synchronize()
        start = time.time()

        # 注意力计算
        scores = torch.matmul(q, k.transpose(-2, -1)) / np.sqrt(d_model // num_heads)
        attn_weights = torch.softmax(scores, dim=-1)
        output = torch.matmul(attn_weights, v)

        torch.cuda.synchronize()
        end = time.time()

        print(f"   序列长度 {seq_len}: {(end - start) * 1000:.2f} ms")


def memory_stress_test():
    """显存压力测试"""
    print("\n=== 显存压力测试 ===")

    device = torch.device("cuda")
    total_memory = torch.cuda.get_device_properties(0).total_memory / 1024 ** 3
    print(f"显卡总显存: {total_memory:.1f} GB")

    # 逐步增加显存使用
    tensors = []
    allocated_gb = 0

    try:
        while allocated_gb < total_memory * 0.9:  # 使用90%显存
            # 每次分配 100MB
            tensor = torch.randn(100 * 1024 * 1024 // 4, device=device)
            tensors.append(tensor)
            allocated_gb += 0.1

            current_allocated = torch.cuda.memory_allocated() / 1024 ** 3
            print(f"\r已分配显存: {current_allocated:.1f} GB", end="", flush=True)

    except RuntimeError as e:
        if "out of memory" in str(e):
            print(f"\n显存不足,最大可用: {torch.cuda.memory_allocated() / 1024 ** 3:.1f} GB")
        else:
            print(f"\n错误: {e}")

    # 清理显存
    del tensors
    torch.cuda.empty_cache()
    print(f"\n显存已清理")


if __name__ == "__main__":
    if torch.cuda.is_available():
        comprehensive_gpu_test()
        memory_stress_test()
        print(f"\n 性能测试完成!")
    else:
        print(" CUDA 不可用")

预期输出应显示正确版本号、True及CUDA 12.8+版本信息。

D:\Programming\software\conda\envs\hu_3_10\python.exe C:\Users\huruiyang\Desktop\pytest\main.py 
=== GPU 性能测试 ===

测试设备: NVIDIA GeForce RTX 5080 Laptop GPU
显存总量: 15.9 GB
CUDA 版本: 12.8
PyTorch 版本: 2.10.0.dev20251119+cu128

1. 矩阵运算性能测试:
   1024x1024: 0.10 ms/次, 21.46 TFLOPS
   2048x2048: 0.70 ms/次, 24.53 TFLOPS
   4096x4096: 7.31 ms/次, 18.80 TFLOPS
   8192x8192: 58.53 ms/次, 18.79 TFLOPS

2. 混合精度 (FP16) 测试:
   FP32: 10.51 ms
   FP16: 122.71 ms
   FP16 加速比: 0.09x

3. 深度学习模拟测试:
   Batch size  1: 132.30 ms, 7.6 imgs/sec
   Batch size  8: 7.51 ms, 1065.0 imgs/sec
   Batch size 16: 10.01 ms, 1598.9 imgs/sec
   Batch size 32: 20.54 ms, 1558.2 imgs/sec

4. 显存带宽测试:
   100 MB: 390.4 GB/s
   500 MB: 390.4 GB/s
   1000 MB: 362.2 GB/s
   2000 MB: 350.4 GB/s

5. Transformer 注意力机制测试:
   序列长度 512: 37.08 ms
   序列长度 1024: 4.50 ms
   序列长度 2048: 342.77 ms

=== 显存压力测试 ===
显卡总显存: 15.9 GB
已分配显存: 14.1 GB
显存已清理

 性能测试完成!

进程已结束,退出代码为 0

Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐