#测试pytorch-gpu是否能用
import torch
flag = torch.cuda.is_available()
print(flag)
ngpu= 1
# Decide which device we want to run on
device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")
print('cuda设备名:',device)
print('gpu名称:',torch.cuda.get_device_name(0))
print('pytorch版本:',torch.__version__)
print('cuda版本:',torch.version.cuda)
print('cudnn版本号:',torch.backends.cudnn.version())
print('定义一个torch格式的3*3的矩阵:',torch.rand(3,3).cuda())
#测试tensorflow-gpu是否能用
import tensorflow as tf
print('\n\nGPU',tf.config.list_physical_devices('GPU'))
a = tf.constant(2.)
b = tf.constant(4.)
print('打印a*b:',a * b)
print("tensorflow版本:", tf.__version__)
正确安装了环境,运行结果如下:
或者用一个简单的官方测试代码来解决测试,此为tensorflow2.0及以上版本使用,初学者的 TensorFlow 2.0 教程 | TensorFlow Core
1.x版本测试可以参考使用前向传播和反向传播的神经网络代码_gaoyueace的博客-CSDN博客_神经网络前向传播和反向传播代码
import tensorflow as tf
mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
model = tf.keras.models.Sequential([
tf.keras.layers.Flatten(input_shape=(28, 28)),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dropout(0.2),
tf.keras.layers.Dense(10, activation='softmax')
])
model.compile(optimizer='adam',loss='sparse_categorical_crossentropy',metrics=['accuracy'])
model.fit(x_train, y_train, epochs=5)
model.evaluate(x_test, y_test, verbose=2)
正确的测试结果如下:
下一条介绍如何正确安装你想要的pytorch-gpu版本和tensorflow-gpu版本。