PyTorch 使用GPU加速

2021-01-02

PyTorch 使用GPU加速

注意：

如果以前安装过CPU版本的PyTorch，务必先卸载，使用pip或conda命令进行卸载。
务必注意PyTorch和CUDA版本之间的对应，比如PyTorch 1.4 请安装CUDA 10.1。
获取CUDA和CuDNN和安装PyTorch GPU版本，以下操作大概率可行。

https://blog.csdn.net/Mind_programmonkey/article/details/99688839#commentBox

检查是否正确安装CUDA，在cmd输入：

nvcc -V

检查是否可用GPU加速，

1
2
3

import torch
print(torch.cuda.is_avaliable())
#输出True 表示成功。

按照以上操作还不行怎么办？
- 尝试去NVIDIA官网更新自己的显卡驱动
- 重启一下电脑？
- 实在不行，在下也没有办法了

使用CUDA加速训练模型

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

import torchvision
import torchvision.transforms as transforms

#定义一个卷积网络对FashionMNIST数据集进行分类
class Network(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5)
        self.conv2 = nn.Conv2d(in_channels=6, out_channels=12, kernel_size=5)
        
        self.fc1 = nn.Linear(in_features=12 * 4 * 4, out_features=120)
        self.fc2 = nn.Linear(in_features=120, out_features=60)
        self.out = nn.Linear(in_features=60, out_features=10)
        
    def forward(self, t):
        # (1) input layer
        t = t
        # (2) hidden conv layer
        t = self.conv1(t)
        t = F.relu(t)
        t = F.max_pool2d(t, kernel_size=2, stride=2)

        # (3) hidden conv layer
        t = self.conv2(t)
        t = F.relu(t)
        t = F.max_pool2d(t, kernel_size=2, stride=2)

        # (4) hidden linear layer
        t = t.reshape(-1, 12 * 4 * 4)
        t = self.fc1(t)
        t = F.relu(t)

        # (5) hidden linear layer
        t = self.fc2(t)
        t = F.relu(t)

        # (6) output layer
        t = self.out(t)
        #t = F.softmax(t, dim=1)
        return t

#训练数据集的加载
train_set = torchvision.datasets.FashionMNIST(
    root='./data'
    ,train=True
    ,download=True
    ,transform=transforms.Compose([
        transforms.ToTensor()
    ])
)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=100,shuffle=True)

if __name__ == "__main__":
    #指定gpu进行训练
    device = torch.device('cuda:0')
    network = Network().to(device)
    optimizer = optim.Adam(network.parameters(), lr=0.01)
    start = time.clock()
    for epoch in range(10):

        total_loss = 0
        total_correct = 0

        for batch in train_loader: # Get Batch
            images, labels = batch 
             #moving the inputs to gpu type
            images,labels = images.to(device), labels.to(device)
            network.eval()
            preds = network(images) # Pass Batch
            loss = F.cross_entropy(preds, labels) # Calculate Loss

            optimizer.zero_grad()
            loss.backward() # Calculate Gradients
            optimizer.step() # Update Weights

            total_loss += loss.item()
            total_correct += get_num_correct(preds, labels)

        print(
            "epoch", epoch, 
            "total_correct:", total_correct, 
            "loss:", total_loss
        )
elapsed = (time.clock() - start)
print(elapsed)

可以发现，使用CUDA比使用CPU训练快多了。