PyTorch 中的 In-place 操作

欢迎关注 “小白玩转Python”，发现更多 “有趣”

今天的深度神经网络有数以百万计的可训练参数，因此可以试图在免费的 GPU 上训练它们，如 Kaggle 或 Google Colab，常常会导致 GPU 上的显存耗尽。有几种简单的方法可以减少模型占用的 GPU 显存，例如：

• 考虑更改模型的体系结构或使用具有更少可训练参数的模型类型，但是这种方法可能会影响模型的性能指标；

• 减少批量大小或手动设置数据加载的数量。在这种情况下，模型需要更长的时间来训练。

在神经网络中使用 in-place 操作可以有助于避免上述方法的缺点，同时节省一些 GPU 显存。但是，由于以下几个原因，不建议使用 in-place 操作。

在这篇文章中，我想：

• 描述什么是 in-place 操作，并演示他们如何可能有助于节省 GPU 的显存；

• 告诉我们为什么要避免 in-place 操作，或者非常谨慎地使用它们。

in-place 操作

in-place 操作是直接改变给定线性代数、向量、矩阵(张量)的内容而不需要复制的运算。

——摘自 Python 教程

根据定义，in-place 操作不复制输入。这就是为什么在处理高维数据时，它们可以帮助减少显存的使用。

我想演示如何就地操作有助于消耗较少的 GPU 显存。为了做到这一点，我将使用这个简单的函数来测试PyTorch分别为 out-of-place 的 ReLU 和 in-place 的 ReLU 分配的显存：

# 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-of-place 的 ReLU 分配的显存：

# 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))

我收到的输出是这样的：

Allocated memory: 382.0
Allocated max memory: 382.0

然后按如下方式来测量 in-place 的 ReLU 分配的显存：

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))

我收到的输出是这样的：

Allocated memory: 0.0
Allocated max memory: 0.0

看起来使用 in-place 操作可以帮助我们节省一些 GPU 显存。但是，在使用 in-place 操作时要格外小心，并进行两次检查。在下一节中，我将告诉你们为什么。

in-place 操作的缺点

它们可能会覆盖计算梯度所需的值，这意味着打破了模型的训练过程。这是 PyTorch 的官方文档所说的：

在 autograd 中支持 in-place 操作是一件困难的事情，我们在大多数情况下都不鼓励使用它们。Autograd 积极的缓冲区释放和重用使其非常高效，而且很少有 in-place 操作真正降低任何显著的显存使用量的情况。除非是在显存压力很大的情况下操作，否则你可能永远都不需要使用它们。

有两个主要原因限制了 In-place 操作的适用性：

• in-place 操作可能会覆盖计算梯度所需的值。

• 每个 in-place 操作实际上都需要重写计算图的实现。out-of-place只是分配新对象并保留对旧计算图的引用，而 in-place 操作则需要将所有输入的创建更改为代表此操作的函数。

小心使用in-place 操作的另一个原因是它们的实现非常棘手。这就是为什么我建议使用 PyTorch 标准的in-place 操作(如上面的in-place  ReLU) ，而不是手动实现。

让我们来看一个 SiLU (或 Swish-1)激活函数的例子，这是 SiLU 的不合适的实现：

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 函数实现 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 _ in-place 实现 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-place 操作及其目的，演示了 in-place 操作如何有助于减少 GPU 显存消耗；

• 我描述了 in-place 操作的显著缺点，使用时要非常小心。

·  END  ·

HAPPY LIFE