正文概述

导读

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分配的内存，使用这个简单的函数：

# 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

# 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分配的内存：

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

# 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

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

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

1 2	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实现：

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

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操作的显著缺点。人们应该非常小心地使用它们，并检查两次结果。

英文原文：https://towardsdatascience.com/in-place-operations-in-pytorch-f91d493e970e

打赏赞

pytorch 模型激活函数神经网络

本站上原创文章未经作者许可，不得用于商业用途，仅做学习交流使用，本站免责声明。转载请注明出处，否则保留追究法律责任的权利。《署名-非商业性使用-相同方式共享 4.0 国际 (CC BY-NC-SA 4.0)》许可协议授权
数据科学与编程 » PyTorch中的In-place操作是什么？为什么要避免使用这种操作？

分享到：