天天看点

PyTorch学习笔记(五)使用GPU训练模型0. 写在前面1. 设置 GPU2. 将数据和模型在 GPU 和 CPU 之间移动3. 多 GPU 数据并行

Environment

  • OS: Ubuntu 14.04
  • Python version: 3.7
  • PyTorch version: 1.4.0
  • IDE: PyCharm
  • GPU: 3 张 RTX 2080 Ti

文章目录

  • 0. 写在前面
  • 1. 设置 GPU
    • 1.1 选择 GPU
    • 1.2 查看可用 GPU 的信息
  • 2. 将数据和模型在 GPU 和 CPU 之间移动
    • 1.2 对于 torch.Tensor
    • 2.2 对于 torch.nn.Module
  • 3. 多 GPU 数据并行

0. 写在前面

深度学习的神经网络模型往往参数巨多,一个能跑的 GPU 是基本配置。这里记一下 PyTorch 中关于 GPU 的一些函数和训练代码。

CPU 和 GPU 的区别可以简明地参考知乎 CPU 和 GPU 的区别是什么?。

1. 设置 GPU

1.1 选择 GPU

以三块 RTX 2080 Ti 卡为例,它们的物理索引分别为 0、1、2

os.system('nvidia-smi -q -d Memory | grep -A4 GPU | grep Free > temp.txt')
gpu_memories = [int(x.split()[2]) for x in open('temp.txt', 'r').readlines()]
os.system('rm temp.txt')
print('GPU of free memories:', gpu_memories)
# GPU of free memories: [11009, 11009, 11009]
           

接着,设置索引为 1 和 2 的两块卡对当前 python 脚本程序可见

这样,物理 GPU 中索引为 1 和 2 的 GPU 为对应索引为 0 和 1 逻辑 GPU。

1.2 查看可用 GPU 的信息

torch.cuda

模块中提供了 PyTorch 关于使用 GPU 的函数。

  • torch.cuda.is_available()

    返回是否有可用的 GPU
import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)  # cuda
           
  • torch.cuda.get_device_name()

    返回 GPU 设备的名称
import torch

print(torch.cuda.get_device_name())  # GeForce RTX 2080 Ti
           
  • torch.cuda.device_count()

    返回当前程序可见的 GPU 数目
import torch

print(torch.cuda.device_count())  # 2
           

2. 将数据和模型在 GPU 和 CPU 之间移动

使用

Tensor.is_cuda()

方法能够查看数据是否在 GPU 上

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tensor_in_cpu = torch.randn((128, 3, 299, 299))
print(tensor_in_cpu.is_cuda)
tensor_in_gpu = tensor_in_cpu.to(device)  # 移动数据到 cuda 这块 GPU
print(tensor_in_gpu.is_cuda)
           

1.2 对于 torch.Tensor

Tensor.to()

方法,转换数据类型或所在设备(CPU / GPU)

  • 传入设备,如

    torch.device(cuda:0)

    ,则将张量移动到索引为 0 的逻辑 GPU 上
  • 传入数据类型,如

    torch.float32

    ,则将张量的类型转换为该类型

注意,该操作非 in-place,需要创建一个变量来引用得到的结果。

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tensor_in_cpu = torch.randn((128, 3, 299, 299))
tensor_in_gpu = tensor_in_cpu.to(device)  # 移动数据到 cuda:0 这块 GPU
           

使用

Tensor.is_cuda()

方法能够查看数据是否在 GPU 上

print(tensor_in_cpu.is_cuda)  # False
print(tensor_in_gpu.is_cuda)  # True
           

2.2 对于 torch.nn.Module

Module.to()

方法,转换模型参数的数据类型或模型所在的设备(CPU / GPU)。

  • 传入设备,如

    torch.device(cuda:0)

    ,则将模型移动到索引为 0 的逻辑 GPU 上
  • 传入数据类型,如

    torch.float32

    ,则将模型中参数的类型转换为该类型

注意,该操作 in-place。

import torchvision

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = torchvision.models.resnet50()
model.to(device)

           

3. 多 GPU 数据并行

更多详细的参考知乎 Pytorch的nn.DataParallel。

torch.nn.DataParallel

类,将模型中的数据分发到不同 GPU 上,让模型水平上的数据并行,实现多 GPU 训练。

import torch
from torch.nn import DataParallel, Module, Sequential, Linear, ReLU

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


class FooNet(Module):
    def __init__(self, in_features, num_classes):
        super(FooNet, self).__init__()
        self.fc = Sequential(
            Linear(in_features, 100),
            ReLU(),
            Linear(100, num_classes)
        )

    def forward(self, x):
        print("Batch size of data in 'forward' method:", x.size(0))
        return self.fc(x)


# ========== test ==========
batch_size = 128
in_features = 100
num_classes = 100

# create data
inputs = torch.randn(batch_size, in_features)
labels = torch.randn(batch_size, num_classes)
inputs, labels = inputs.to(device), labels.to(device)

# modeling
foonet = FooNet(in_features, num_classes)
foonet = DataParallel(
    module=foonet,  # 需要包装分发的模型
    device_ids=None,  # 可分发的 GPU,默认 None,为分发到所有可见可用的 GPU
    output_device=None,  # 结果输出的设备,默认 None,为主 GPU,即索引为 0 的逻辑 GPU
    dim=0  # 应该是按维度 0 将数据分发到不同 GPU 上(文档并没有解释)
)
foonet.to(device)

# forward
outputs = foonet(inputs)
print('Size of output:', outputs.size())

           

运行结果为打印出

Batch size of data in 'forward' method: 64
Batch size of data in 'forward' method: 64
Size of output: torch.Size([128, 100])
           

可见,输入数据一个 batch 为 128,被平均分配到两张卡中,每张卡的模型中数据 batch 为 64。

注意,

DataParallel

包装后的模型是

torch.nn.parallel.data_parallel.DataParallel

类的实例对象,需要用

.module

属性得到具体的模型

print(type(foonet))
# <class 'torch.nn.parallel.data_parallel.DataParallel'>

print(type(foonet.module))
# <class '__main__.FooNet'>
           

继续阅读