CPU与GPU上检测pytorch是否安装成功

python学习

python学习

0.安装pytorch

官方安装pytorch地址

1.验证pytorch已经安装成功

1.1确定pytorch版本

import torch
print(torch.__version__)

如果没有报错，同时输出了正确版本号，就说明PyTorch已经成功安装了。

1.2.测试pytorch基础功能

为了确保PyTorch基础功能正常工作，可以进行以下测试：
1.测试是否可以创建张量：

import torch
x = torch.Tensor(5, 3)
print(x)

如果没有报错，同时输出了5行3列的张量，就说明可以创建张量。
2.测试是否可以进行张量计算

import torch
x = torch.Tensor(5, 3)
y = torch.Tensor(3, 4)
z = torch.mm(x, y)
print(z)

如果没有报错，同时输出了5行4列的张量，就说明可以进行张量计算。

1.3.在GPU上测试pytorch

如果电脑有GPU支持，建议在GPU上测试PyTorch是否正常工作。测试代码如下：

import torch
if torch.cuda.is_available():
    x = torch.Tensor(5, 3).cuda()
    y = torch.Tensor(3, 4).cuda()
    z = torch.mm(x, y)
    print(z)

如果没有报错，同时输出了5行4列的张量，就说明PyTorch GPU版本也正常工作。

1.4使用实例代码测试

为了更加确认PyTorch是否正常工作，可以使用PyTorch官方提供的示例代码进行测试，例如以下代码：

import torch
import torchvision
model = torchvision.models.resnet18(pretrained=True)
model.eval()
x = torch.rand(1, 3, 224, 224)
y = model(x)

print(y)

如果没有报错，同时输出了正确的结果，就说明PyTorch正常工作。