简介

    gpu(图形处理器),全称Graphics Processing Unit。GPU虽然只能做简单的运算,但其核心数多以及应用并行运算,非常适合矩阵运算。CPU单核运算能力强大,但也比不过GPU的上千个核心的运算能力。深度学习的模型训练,运用的是矩阵运算,因此,GPU天生适合神经网络的数据处理。

  • 放cpu和gpu训练时的表现,就可以看到gpu的魅力所在了
    cpu

    gpu
  • Pytorch(torch)使用GPU的配置

    1. 查看cuda版本

    nvidia-smi
    

    2. 驱动下载
    https://developer.nvidia.com/rdp/cudnn-archive#a-collapse742-10

  • 第一次下载需要注册为NVIDIA会员
  • 3. python安装pytorch、cudatoolkit

  • 注意pytorch和cuda版本对应
    https://pytorch.org/get-started/previous-versions/
  • 如何使用GPU

    1. 查看能否使用gpu
    import torch
    torch.cuda.is_available() 
    

    返回True,说明可用;如果返回False,可能是你的pytorch和cuda版本不一致造成的。

    1. 看官网的创建神经网络的案例领悟

    https://pytorch.org/tutorials/beginner/basics/buildmodel_tutorial.html

    import os
    import torch
    from torch import nn
    from torch.utils.data import DataLoader
    from torchvision import datasets, transforms
    
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using {device} device")
    
    class NeuralNetwork(nn.Module):
        def __init__(self):
            super(NeuralNetwork, self).__init__()
            self.flatten = nn.Flatten()
            self.linear_relu_stack = nn.Sequential(
                nn.Linear(28*28, 512),
                nn.ReLU(),
                nn.Linear(512, 512),
                nn.ReLU(),
                nn.Linear(512, 10),
            )
    
        def forward(self, x):
            x = self.flatten(x)
            logits = self.linear_relu_stack(x)
            return logits
    
    model = NeuralNetwork().to(device) #在这里。。。
    print(model)
    

    基本上都是用to函数来将计算放到gpu上。或者也可以在创建对象的时候指定设备,再看例子。

    X = torch.rand(1, 28, 28, device=device)
    logits = model(X)
    pred_probab = nn.Softmax(dim=1)(logits)
    y_pred = pred_probab.argmax(1)
    print(f"Predicted class: {y_pred}")
    

    来源:维度攻城狮

    物联沃分享整理
    物联沃-IOTWORD物联网 » torch使用gpu的环境配置

    发表评论