博客

Linux安装Pytorch1.8GPU(CUDA11.1)的实现

  • Post category:Python

下面是详细讲解“Linux安装Pytorch1.8GPU(CUDA11.1)的实现”的完整攻略。在进行安装之前，需要确保计算机上已经安装好了相应的CUDA和cuDNN。

第一步：安装Anaconda

可以在Anaconda官网上下载对应的安装程序，选择相应的操作系统和版本即可。

第二步：创建虚拟环境

在终端中输入以下命令来创建一个名为“pytorch”（也可以自己取一个名字）的虚拟环境：

conda create --name pytorch python=3.8

第三步：激活虚拟环境并安装PyTorch

在终端中输入以下命令激活虚拟环境：

conda activate pytorch

接着，使用pip安装PyTorch：

pip install torch torchvision torchaudio -f https://download.pytorch.org/whl/cu111/torch_stable.html

-f参数是指在指定的URL中寻找stable版的PyTorch。

第四步：测试安装

输入以下命令以测试PyTorch是否成功安装：

import torch

x = torch.randn(2, 3).cuda()
print(x)

如果没有报错，输出结果应该是类似于以下的内容：

tensor([[-0.6672,  1.1235, -0.8539],
        [-0.2701, -1.2901, -2.0667]], device='cuda:0')

示例1：在代码中使用PyTorch

以下演示如何使用安装好的PyTorch来构建一个简单的神经网络，代码示例如下：

import torch
import torch.nn as nn
import torch.optim as optim

# 构建神经网络
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(784, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, 10)

    def forward(self, x):
        x = x.view(-1, 784)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 训练模型
net = Net().cuda()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.01)

for epoch in range(10):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        inputs, labels = inputs.cuda(), labels.cuda()

        optimizer.zero_grad()

        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    print('[%d, %5d] loss: %.3f' %
          (epoch + 1, i + 1, running_loss / 1000))
    running_loss = 0.0

示例2：使用CUDA加速矩阵计算

以下演示如何使用安装好的PyTorch和CUDA来加速矩阵计算，代码示例如下：

import torch
import time

# 定义矩阵大小
batch_size, matrix_size = 128, 8192

# 创建CPU和GPU上的随机矩阵
x_cpu, y_cpu = torch.randn(batch_size, matrix_size), torch.randn(batch_size, matrix_size)
x_gpu, y_gpu = x_cpu.cuda(), y_cpu.cuda()

# 用CPU计算矩阵乘法
start = time.time()
res_cpu = x_cpu.mm(y_cpu.t())
end = time.time()
print("CPU: ", end - start)

# 用GPU计算矩阵乘法
start = time.time()
res_gpu = x_gpu.mm(y_gpu.t())
end = time.time()
print("GPU: ", end - start)

输出结果应该类似于以下内容：

CPU:  3.500962495803833
GPU:  0.006424903869628906

可以看到，使用了GPU后，矩阵计算的速度大幅提升。

Tags: