[CUDA] [Ubuntu] 安装 Tensorflow/Pytorch

Conda 安装教程

先用conda复制一个虚拟环境

conda create --name vm_jupyter_tensorflow --clone vm_jupyter

激活环境

conda activate vm_jupyter_tensorflow

查看cuda要求, 记得选英文, 按照这个教程 安装对应版本的cuda和cuDNN

安装tensorflow

pip install tensorflow

打开notebook, 检测一下tensorflow是否正常调用

发现无法正常调用, 原来要重启一下, 重启之后可以看到正常调用了

import tensorflow as tf

gpus = tf.config.list_physical_devices('GPU')
if gpus:
    print(f"Available GPU(s): {gpus}")
else:
    print("No GPU detected")

然后再来安装pytorch. 就不创建新的虚拟环境了, 步骤一样的

首先去官网查看一下cuda要求

我的cuda版本不够, 先装个12.4的cuda, 如何切换cuda和cuDNN版本

复制指令并安装

pip3 install torch torchvision torchaudio

启动notebook检查一下, 已经完美识别

import torch

if torch.cuda.is_available():
    print("GPU is available")
    # 检查当前的 GPU 数量
    print(f"Number of GPUs available: {torch.cuda.device_count()}")
    # 打印当前使用的 GPU 名称
    print(f"Current GPU: {torch.cuda.get_device_name(torch.cuda.current_device())}")
else:
    print("No GPU available, using CPU instead.")