导读
In-place操作用在推理的时候可以显著节省内存，但是训练的时候一定要小心使用。
如今的高级深度神经网络拥有数百万个可训练参数，训练它们通常会导致GPU内存耗尽。有几种简单的方法可以减少模型占用的GPU内存，例如：
- 考虑改变模型的架构或使用具有较少可训练参数的模型类型(例如，选择DenseNet-121而不是DenseNet-169)。这种方法会影响模型的性能度量。
- 减少batch大小或手动设置数据加载workers的数量。在这种情况下，模型需要更长的时间来训练。
在神经网络中使用in-place操作可能有助于避免上述方法的缺点，同时节省一些GPU内存。但是，由于几个原因，不建议使用in-place操作。
在这篇文章中，内容包括：
- 描述什么是in-place操作，并演示他们如何可能有助于节省GPU内存。
- 告诉我们为什么要避免in-place操作或非常小心地使用它们。
In-place 操作
“In-place运算是一种直接改变给定线性函数、向量、矩阵(张量)内容而不复制的运算。"
根据定义，in-place操作不会复制输入。这就是为什么它们可以帮助在操作高维数据时减少内存使用。
我想演示in-place操作如何帮助消耗更少的GPU内存。为了做到这一点，我将度量out- place ReLU和in-place ReLU分配的内存，使用这个简单的函数：
代码语言：javascript复制# Import PyTorch
import torch # import main library
import torch.nn as nn # import modules like nn.ReLU()
import torch.nn.functional as F # import torch functions like F.relu() and F.relu_()
def get_memory_allocated(device, inplace = False):
'''
Function measures allocated memory before and after the ReLU function call.
INPUT:
- device: gpu device to run the operation
- inplace: True - to run ReLU in-place, False - for normal ReLU call
'''
# Create a large tensor
t = torch.randn(10000, 10000, device=device)
# Measure allocated memory
torch.cuda.synchronize()
start_max_memory = torch.cuda.max_memory_allocated() / 1024**2
start_memory = torch.cuda.memory_allocated() / 1024**2
# Call in-place or normal ReLU
if inplace:
F.relu_(t)
else:
output = F.relu(t)
# Measure allocated memory after the call
torch.cuda.synchronize()
end_max_memory = torch.cuda.max_memory_allocated() / 1024**2
end_memory = torch.cuda.memory_allocated() / 1024**2
# Return amount of memory allocated for ReLU call
return end_memory - start_memory, end_max_memory - start_max_memory
调用该函数来测量out- place ReLU分配的内存：
代码语言：javascript复制# setup the device
device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu")
# call the function to measure allocated memory
memory_allocated, max_memory_allocated = get_memory_allocated(device, inplace = False)
print('Allocated memory: {}'.format(memory_allocated))
print('Allocated max memory: {}'.format(max_memory_allocated))
得到结果如下：
代码语言：javascript复制Allocated memory: 382.0
Allocated max memory: 382.0
然后调用in-place ReLU：
代码语言：javascript复制memory_allocated_inplace, max_memory_allocated_inplace = get_memory_allocated(device, inplace = True)
print('Allocated memory: {}'.format(memory_allocated_inplace))
print('Allocated max memory: {}'.format(max_memory_allocated_inplace))
得到结果如下：
代码语言：javascript复制Allocated memory: 0.0
Allocated max memory: 0.0
看来使用in-place操作可以帮助我们节省一些GPU内存。然而，我们在使用现场操作时应该非常谨慎，并且要反复检查。在接下来的部分，我将告诉你为什么。
In-place 操作的缺点
in-place操作的主要缺点是，它们可能会覆盖计算梯度所需的值，这意味着破坏模型的训练过程。这是PyTorch autograd官方文档所说的：
在autograd支持in-place操作是一件困难的事情，我们在大多数情况下不鼓励使用它们。Autograd的主动缓冲区释放和重用使其非常高效，在很少情况下，in-place操作实际上会显著降低内存使用量。除非你正在承受巨大的内存压力，否则你可能永远都不需要使用它们。 限制in-place作业的适用性的主要原因有两个： 1、in-place操作可能会覆盖计算梯度所需的值。 2、每个in-place操作实际上都需要实现重写计算图。Out-of-place版本只是简单地分配新对象并保持对旧图的引用，而in-place操作则要求将所有输入的创建者更改为表示该操作的函数。
要小心使用in-place操作的另一个原因是，它们的实现非常棘手。这就是为什么我建议使用PyTorch标准的in-place操作(如上面的就地ReLU)，而不是手动实现。
让我们看一个SiLU(或Swish-1)激活函数的例子。以下是SiLU的out-of-place实现：
代码语言：javascript复制def silu(input):
'''
Out-of-place implementation of SiLU activation function
https://arxiv.org/pdf/1606.08415.pdf
'''
return input * torch.sigmoid(input)
我们尝试使用torch.sigmoid_实现一个 in-place SiLU：
def silu_inplace_1(input):
'''
Incorrect implementation of in-place SiLU activation function
https://arxiv.org/pdf/1606.08415.pdf
'''
return input * torch.sigmoid_(input) # THIS IS INCORRECT!!!
上面的代码不正确地实现了in-place SiLU。只要比较两个函数返回的值，就可以确定。实际上，函数silu_inplace_1 返回sigmoid(input) * sigmoid(input)！使用torch.sigmoid_实现SiLU的例子：
def silu_inplace_2(input):
'''
Example of implementation of in-place SiLU activation function using torch.sigmoid_
https://arxiv.org/pdf/1606.08415.pdf
'''
result = input.clone()
torch.sigmoid_(input)
input *= result
return input
这个小示例演示了为什么我们在使用in-place操作时应该谨慎并检查两次。
总结
本文简介:
- 说明了in-plac及其目的。演示了in-plac操作如何帮助消耗更少的GPU内存。
- 描述了in-plac操作的显著缺点。人们应该非常小心地使用它们，并检查两次结果。
