pip安装tensorflow-gpu2.3.0和pytorch-gpu1.7.1

前提条件

需要安装好CUDA10.1和CUDNN7.6.

安装tensorflow-gpu

查看对应版本

https://tensorflow.google.cn/install/source_windows#gpu

如上图，CUDA10.1和CUDNN7.6环境对应的tensorflow-gpu版本有2.2.0和2.3.0，以2.3.0为例。

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple tensorflow-gpu==2.3.0

测试代码

import tensorflow as tf
import os
os.environ['TF_CPP_MIN_LOG_LEVEL']='2'
print(tf.__version__)
a = tf.constant(1.)
b = tf.constant(2.)
print(a+b)
print('GPU:', tf.test.is_gpu_available())

输出结果

安装pytorch-gpu

打开官网

https://pytorch.org/get-started/locally/

pip install torch==1.7.1+cu101 torchvision==0.8.2+cu101 torchaudio===0.7.2 -f https://download.pytorch.org/whl/torch_stable.html

测试代码

import torch
print(torch.__version__)
print(torch.cuda.is_available())

输出结果