PyTorch 101，Part4：内存管理以及使用多个GPU
这篇文章涉及到PyTorch高级的GPU管理，包括如何为你的神经网络配置多个GPU，是否使用数据或者模型的并行计算。我们用最佳的实践总结了如何调试内存错误。
这里是我们PyTorch 101系列的第四部分，在这篇文章中，我们将会涉及到多GPU的使用。
在这个部分我们将会涉及：
- 如何让你的网络使用多个GPU，使用数据或者模型并行计算。
- 当创建一个新的对象的时候，如何自动地选择GPU
- 如何诊断和分析出现的内存问题
目录
PyTorch 101，Part4：内存管理以及使用多个GPU
1 在CPU和GPUs之间移动张量
在PyTorch中，每个张量都有一个to()的成员变量。它的任务是将张量放在一个具体设备上，可能是CPU或者是GPU。to函数的输入是一个torch.device的实例对象，它可以用下面的输入进行初始化：
- cpu:将数据放在CPU上
- cuda:0 将数据放在第0个GPU上。同样，你可以把张量放置在别的GPU上面。
一般来说，无论什么时候你初始化一个Tensor，它都会默认放在CPU上面。你之后可以将它移动到GPU上面。你可以通过调用torch.cuda.is_available函数来检查GPU是否可用。
if torch.cuda.is_available():
dev = "cuda:0"
else:
dev = "cpu"
device = torch.device(dev)
# 或者使用
# device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
a = torch.zeros(4,3)
a = a.to(device) #alternatively, a.to(0)
你也可以通过传入一个索引作为to函数的参数，将张量移动到某个具体的GPU上面。重要的是，上面的代码和设备无关，也就是说，你不需要单独改变它，就可以让数据既能在CPU上运行又能在GPU上运行（这完全根据你电脑的配置来定）。
1.1 cuda()函数
另一个将张量放置在GPU上的方法是调用cuda(n)函数，其中n代表的是GPU的索引。如果你调用了cuda，那么这个张量就会放置在CPU n上面。
torch.nn.Module类还有to函数和cuda函数，它可以将整个网络放在一个特定的设备上面。Tensor在nn.Module的实例对象上调用to函数就够了，并不需要分配to函数的返回值。
clf = myNetwork()
clf.to(torch.device("cuda:0") # or clf = clf.cuda()
# 或者使用：clf = myNetwork().to(device)
2 自动选择GPU
尽管可以显式地定义张量在哪个GPU上运行，但是在进行运算的过程中，我们会创建很多张量，我们希望它能够自动地在某个设备上进行创建，这样能够减少张量在设备间进行交换，这会降低代码执行的速度。在这方面，PyTorch给我们提供了一些函数去实现这个思路。
第一个函数是torch.get_device()函数。它只支持GPU类型的张量，这个函数的返回值张量所在GPU的索引。我们可以使用这个函数取决定张量的所在的设备，因此我们可以将一个创建好了的张量自动地移动到这个设备上。
#making sure t2 is on the same device as t2
a = t1.get_device()
b = torch.tensor(a.shape).to(dev)
当我们创建一个新的Tensors的时候，我们也可以调用cuda(n)函数。通过调用cuda函数所创建的张量默认放置在GPU0上面，但是通过下面的声明对其进行改变。
torch.cuda.set_device(0) # or 1,2,3
如果在两个操作数之间创建了一个张量，这两个操作数在相同的设备上，那么这个张量也会放置在这个设备上。但是如果两个操作数在不同的设备上上面，这将会报错。
2.1 new_*函数
在PyTorch1.0版本中，还可以利用new_函数实现这样的结果。当一个Tensor调用了形如new_ones函数的时候，它会返回一个和这个Tensor相同的数据类型，以及将其放置在调用new_ones张量相同的设备上面。
ones = torch.ones((2,)).cuda(0)
# Create a tensor of ones of size (3,4) on same device as of "ones"
newOnes = ones.new_ones((3,4))
randTensor = torch.randn(2,4)
更多new_函数可也在PyTorch文档中找到，相关连接我放在下面了。
3 使用多个GPUs
我们有两个使用多个GPUs的方法：
- 数据并行计算。我们可以将一个大的batches分成很多小一点的batches，并且在多个GPUs上面并行处理这些小的batches
- 模型并行计算。我们可以将一个神经网络分成几个小一点的子网络，然后再不同的GPU上面处理这些子网络。
3.1 数据并行计算
在PyTorch中中，可以使用nn.DataParallel类来实现数据并行计算。你可以用一个表示你的神经网络的nn.Module实例对象来初始化一个nn.DataParallel实例对象，然后传入一连串的GPU索引号，这些表示batches要在哪些GPU上进行处理。
parallel_net = nn.DataParallel(myNet, gpu_ids = [0,1,2])
现在，你可以简单地执行nn.DataParallel对象，就像使用一个nn.Module对象一样。
predictions = parallel_net(inputs) # Forward pass on multi-GPUs
loss = loss_function(predictions, labels) # Compute loss function
loss.mean().backward() # Average GPU-losses + backward pass
optimizer.step()
但是，我需要说清楚一些事情。尽管我们的数据在多个GPUs上并行计算，但是在一开始我们需要将数据保存在一个GPU上面。
我们还需要确保DataParallel实例对象也在特定的GPU上面。这个语法和我们之前在使用nn.Module的相似。
input = input.to(0)
parallel_net = parellel_net.to(0)
事实上，下面的示例图表述了nn.DataParallel是如何工作的。
DataParallel接收输入，将数据分成多个更小的batches，在所有的GPU上覆制神经网络，执行前向传播，然后再原来的GPU上整合输出。
DataParallel存在一个问题，那就是在一个GPU上面（就是主节点）会不均衡的加载数据。通常来说有两个方法可以解决这些问题。
- 首先，在前向传播的时候计算损失。这可以确保至少是并行计算网络的损失函数。
- 另一个方法是去实现一个并行的损失函数网络层。这个超过这这篇文章的范围，但是，如果感兴趣的话，我给出了相关文章的链接，他们详细介绍了如何实现这样的网络层。
3.2 模型并行计算
模型并行计算意味着你可以将你的网络模型划分成多个更小的子网络，然后将他们放在不同的GPU上面。这样做的主要是因为如果你的网络模型太大，一个GPU放不下。
注意，模型并行计算比数据并行计算要慢，这是因为将一个网络模型放在多个GPUs上面会增加GPUs之间的以来，这样会妨碍它们完全地进行并行计算。这样并不是出于速度快的优点，而是因为它可以在网络模型尺寸太大，一个GPU放不下的时候去运行网络模型。
正如图b所示，在前向传播的时候，子网络2在等待子网络1运行完成，而在反向传播的时候，子网络1要等到子网络2运行完成。
只要你记住两件事，在PyTorch实现模型并行计算就会非常的容易。
- 输入数据和网络模型应该在相同的设备上面。
- to()函数和cuda()函数都支持autograd，因此在反向传播的过程中，你可以把梯度从一个GPU复制到另外一个GPU上面。
我们使用下面的代码来更好的理解。
class model_parallel(nn.Module):
def __init__(self):
super().__init__()
self.sub_network1 = ...
self.sub_network2 = ...
self.sub_network1.cuda(0)
self.sub_network2.cuda(1)
def forward(x):
x = x.cuda(0)
x = self.sub_network1(x)
x = x.cuda(1)
x = self.sub_network2(x)
return x
在init()函数中，我们单独将子网络放在GPU0和1上面。
注意，在forward()函数中，在将中间数据传给sub_network2之前，我们将其从sub_network1转为了GPU 1。因为cuda支持autograd，为了接下来的反向传播，sub_network2的损失的反向传播会复制到sub_network1的缓存器上。
4 解决内存溢出(OOM)错误
在这部分中，我们会涉及如何诊断内存问题，以及如果你的网络超出它需要的内存的可能的解决方法。
尽管内存溢出可能需要减少批次的大小，但是最佳的手段是去检查内存的使用。
4.1 使用GPUtil去跟踪内存的使用情况
跟踪GPU使用的方法是使用nvidia-smi命令行控制台上监控内存的使用情况。这个方法有个问题就是会泄露GPU的使用情况，并且内存泄露发生的太快，你无法准确定位你的代码的哪个部分导致了这个错误。
为了解决这个问题，我们可以使用GPUtil扩展模块，你可以运行下面的命令，使用pip进行安装。
pip install GPUtil
它的使用也非常的简单
import GPUtil
GPUtil.showUtilization()
这需要把第二行代码放在你想检查GPU使用情况的地方就行。通过将这行代码放在代码中的不同位置，你就可以检查究竟是哪个部分导致了你的网络模型的OOM错误。
让我们讨论一下解决OOM错误可能的方法
4.2 使用del关键字处理内存丢失
PyTorch有着非常积极的垃圾回收机制。只要变量超出了范围，垃圾回收就会释放它。
我们需要记住，Python不会像其他语言比如c/c++那样强制执行上面的规则。只要一个变量没有其他指针指向它的时候，它就会被释放（这也就是变量在Python中不需要声明的原因）。
事实上，你的input，output张量所占用的内存尽管超出了训练范围，但是仍然不会被释放。思考一下下面的代码：
运行上面的代码，仍会输出i的值，尽管已
for x in range(10):
i = x
print(i) # 9 is printed
尽管超出了我们初始化i的循环。同样的，loss和output的张量在训练循环之外依然存在。我们可以使用del关键字来真正地释放这些张量。
del out, loss
作为一般的经验法则，如果你处理了一个张量，你应该del它，因为它不会被当做垃圾进行回收，除非没有指针指向它。
4.3 使用Python数据格式而非一维的张量
在训练循环中，我们通常对一些值进行求和去计算一些矩阵。一个大型的例子就是在每次迭代中，我们需要去更新损失。但是，如果在PyTorch中处理的不小心的话，一些事情可能会导致超出它所需要的内存的使用。
思考下面的代码：
total_loss = 0
for x in range(10):
# assume loss is computed
iter_loss = torch.randn(3,4).mean()
iter_loss.requires_grad = True # losses are supposed to differentiable
total_loss += iter_loss # use total_loss += iter_loss.item) instead
在接下来的迭代中，我们需要iter_loss的引用会指定到一个新的iter_loss上，然后代表之前声明的iter_loss对象会被释放。但是事与愿违，为什么？
因为iter_loss是可微的，total_loss += iter_loss这行代码使用一个AddBakward函数结点构建了一个计算图。在接下来的迭代中，AddBakward结点会添加到这个图中，而且保存iter_loss的对象不会被释放。通常来说，当backward调用的时候，分配给计算图的内存会被释放，但是在这里，无法调用backward函数。
解决这个问题的方法是给total_loss加上一个Python数据，而非一个张量，这可能有效的预防计算图的创建。
我们仅仅使用total_loss+=iter_loss代替total_loss++iter_loss.item()。tiem从一个值保存数值的张量中返回python数据。
4.4 清空Cuda寄存器
尽管PyTorch会积极地释放内存，但是PyTorch进程可能不会将内存返回给操作系统，尽管你已经del了你的张量了。这个内存已经保存了，因此它可以很快的分配给一个新的张量，而不请求操作系统新的额外内存。
当你在工作流程中使用了超过两个进程的时候，将会报错。
当第二个进行运行的时候，尽管第一个进程导致了OOM，但是它还是会占用GPU的内存。为了解决这个问题，你可以拿个在代码的最后加上这行命令。
torch.cuda.empy_cache()
我们需要保证被进程占用的空间已经释放了。
import torch
from GPUtil import showUtilization as gpu_usage
print("Initial GPU Usage")
gpu_usage()
tensorList = []
for x in range(10):
tensorList.append(torch.randn(10000000,10).cuda()) # reduce the size of tensor if you are getting OOM
print("GPU Usage after allcoating a bunch of Tensors")
gpu_usage()
del tensorList
print("GPU Usage after deleting the Tensors")
gpu_usage()
print("GPU Usage after emptying the cache")
torch.cuda.empty_cache()
gpu_usage()
下面的结果是在Tesla K80上执行上面代码的结果：
Initial GPU Usage
| ID | GPU | MEM |
------------------
| 0 | 0% | 5% |
GPU Usage after allcoating a bunch of Tensors
| ID | GPU | MEM |
------------------
| 0 | 3% | 30% |
GPU Usage after deleting the Tensors
| ID | GPU | MEM |
------------------
| 0 | 3% | 30% |
GPU Usage after emptying the cache
| ID | GPU | MEM |
------------------
| 0 | 3% | 5% |
4.5 在推理中使用torch.no_grad()
PyTorch在前向传播的时候默认会构建一个计算图。在创建计算图的过程中，它会分配缓存去存储梯度和中间值，在反向传播的过程中，会使用它们去计算梯度。
在反向传播的过程中，除了为叶子结点分配的缓存器之外，所有的缓存器都是被释放。
但是，在推理的过程中，不存在反向传播，并且这些缓存器从来不会释放，这样会导致内存堆积。因此，当你想去执行的代码不需要进行反向传播的时候，将代码放在torch.no_grad()上下文中。
with torch.no_grad()
# your code
5 总结
在这里总结了关于内存管理和PyTorch中多个GPUs的使用。下面是一些你可能想去跟进的文章的重要链接。
6 扩展阅读
- PyTorch
newfunctions - Parallelised Loss Layer: Training Neural Nets on Larger Batches: Practical Tips for 1-GPU, Multi-GPU & Distributed setups
- GPUtil Github page