torch使用gpu的环境配置

简介

    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

如何使用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}")
  • 1
    点赞
  • 18
    收藏
    觉得还不错? 一键收藏
  • 打赏
    打赏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

维度攻城狮

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值