PyTorch学习笔记（五）使用GPU训练模型

Environment

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

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 脚本程序可见

os.environ['CUDA_VISIBLE_DEVICES'] = '1, 2'

这样，物理 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_devic