为PyTorch提供GPU支持

PyTorch-GPU安装

安装docker及nvidia-docker

在Ubuntu上安装Docker并使得Docker支持GPU

安装PyTorch

在宿主机上安装GPU驱动

• 查找合适的Nvidia驱动器版本并安装

  sudo ubuntu-drivers devices
  sudo ubuntu-drivers autoinstall
  

在Docker安装PyTorch

• 根据驱动满足的最大cuda版本来查找合适版本的pytorch，不需要单独安装cuda

  pip install torch==1.12.1+cu102 torchvision==0.13.1+cu102 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu102
  

  Start Locally | PyTorch

• 测试pytorch

  import torch
  print(torch.__version__)
  print(torch.cuda.is_available())
  for i in range(torch.cuda.device_count()):
  	print(torch.cuda.get_device_name(i))
  

  • tensor默认存储在内存中供CPU使用，通过访问其.device成员变量查看其所处设备
    • 创建tensor时，通过devic 参数指定其运行的设备
    • 调用tensor.to(device)将tensor转移到指定设备
    • 调用tensor.cuda(i)将tensor转移到第 i 块GPU及相应的显存（i从0开始，0可省略不写）
    • 调用tensor.cpu() 将tensor转移到CPU

  import torch
  
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  y = torch.tensor([1, 2, 3], device=device)
  print(y.device)
  # cuda:0
  
  x = torch.tensor([1, 2, 3])
  print(x.device)
  # cpu
  x = x.cuda(0)
  print(x.device)
  # cuda:0
  
  z = torch.tensor([1, 2, 3]).to(device)
  print(z.device)
  # cuda:0
  z = z.cpu()
  print(z.device)
  # cpu
  

  Pytorch入门（5）—— 使用 GPU 进行计算

PyTorch案例

函数拟合

# -*- coding: utf-8 -*-

import torch
import math


dtype = torch.float
#device = torch.device("cpu")
device = torch.device("cuda:0")

# Create random input and output data
x = torch.linspace(-math.pi, math.pi, 2000, device=device, dtype=dtype)
y = torch.sin(x)

# Randomly initialize weights
a = torch.randn((), device=device, dtype=dtype)
b = torch.randn((), device=device, dtype=dtype)
c = torch.randn((), device=device, dtype=dtype)
d = torch.randn((), device=device, dtype=dtype)

learning_rate = 1e-6
for t in range(2000):
    # Forward pass: compute predicted y
    y_pred = a + b * x + c * x ** 2 + d * x ** 3

    # Compute and print loss
    loss = (y_pred - y).pow(2).sum().item()
    if t % 100 == 99:
        print(t, loss)

    # Backprop to compute gradients of a, b, c, d with respect to loss
    grad_y_pred = 2.0 * (y_pred - y)
    grad_a = grad_y_pred.sum()
    grad_b = (grad_y_pred * x).sum()
    grad_c = (grad_y_pred * x ** 2).sum()
    grad_d = (grad_y_pred * x ** 3).sum()

    # Update weights using gradient descent
    a -= learning_rate * grad_a
    b -= learning_rate * grad_b
    c -= learning_rate * grad_c
    d -= learning_rate * grad_d


print(f'Result: y = {a.item()} + {b.item()} x + {c.item()} x^2 + {d.item()} x^3')

Learning PyTorch with Examples