当前位置 博文首页 > PyTorch中的拷贝与就地操作详解

    PyTorch中的拷贝与就地操作详解

    作者:Chris_34 时间:2021-08-05 18:26

    前言

    PyTroch中我们经常使用到Numpy进行数据的处理,然后再转为Tensor,但是关系到数据的更改时我们要注意方法是否是共享地址,这关系到整个网络的更新。本篇就In-palce操作,拷贝操作中的注意点进行总结。

    In-place操作

    pytorch中原地操作的后缀为_,如.add_()或.scatter_(),就地操作是直接更改给定Tensor的内容而不进行复制的操作,即不会为变量分配新的内存。Python操作类似+=或*=也是就地操作。(我加了我自己~)

    为什么in-place操作可以在处理高维数据时可以帮助减少内存使用呢,下面使用一个例子进行说明,定义以下简单函数来测量PyTorch的异位ReLU(out-of-place)和就地ReLU(in-place)分配的内存:

    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
     # setup the device
    device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu")
    #开始测试
    # Call the function to measure the allocated memory for the out-of-place ReLU
    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
    '''
    #Then call the in-place ReLU as follows:
    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
    '''
    

    看起来,使用就地操作可以帮助我们节省一些GPU内存。但是,在使用就地操作时应该格外谨慎。

    就地操作的主要缺点主要原因有2点,官方文档:

    1.可能会覆盖计算梯度所需的值,这意味着破坏了模型的训练过程。

    2.每个就地操作实际上都需要实现来重写计算图。异地操作Out-of-place分配新对象并保留对旧图的引用,而就地操作则需要更改表示此操作的函数的所有输入的创建者。

    在Autograd中支持就地操作很困难,并且在大多数情况下不鼓励使用。Autograd积极的缓冲区释放和重用使其非常高效,就地操作实际上降低内存使用量的情况很少。除非在沉重的内存压力下运行,否则可能永远不需要使用它们。

    总结:Autograd很香了,就地操作要慎用。

    拷贝方法

    浅拷贝方法: 共享 data 的内存地址,数据会同步变化

    * a.numpy() # Tensor—>Numpy array

    * view() #改变tensor的形状,但共享数据内存,不要直接使用id进行判断

    * y = x[:] # 索引

    * torch.from_numpy() # Numpy array—>Tensor

    * torch.detach() # 新的tensor会脱离计算图,不会牵扯梯度计算。

    * model:forward()

    还有很多选择函数也是数据共享内存,如index_select() masked_select() gather()。

    以及后文提到的就地操作in-place。

    深拷贝方法:

    * torch.clone() # 新的tensor会保留在计算图中,参与梯度计算

    下面进行验证,首先验证浅拷贝:

    import torch as t
    import numpy as np
    a = np.ones(4)
    b = t.from_numpy(a) # Numpy->Tensor
    print(a)
    print(b)
    '''输出:
    [1. 1. 1. 1.]
    tensor([1., 1., 1., 1.], dtype=torch.float64)
    '''
    b.add_(1)# add_会修改b自身
    print(a)
    print(b)
    '''输出:
    [2. 2. 2. 2.]
    tensor([2., 2., 2., 2.], dtype=torch.float64)
    b进行add操作后, a,b同步发生了变化
    '''
    

    Tensor和numpy对象共享内存(浅拷贝操作),所以他们之间的转换很快,且会同步变化。

    造torch中y = x + y这样的运算是会新开内存的,然后将y指向新内存。为了进行验证,我们可以使用Python自带的id函数:如果两个实例的ID一致,那么它们所对应的内存地址相同;但需要注意是在torch中还有些特殊,数据共享时直接打印tensor的id仍然会出现不同。

    x = torch.tensor([1, 2])
    y = torch.tensor([3, 4])
    id_0 = id(y)
    y = y + x
    print(id(y) == id_0) 
    # False 
    

    这时使用索引操作不会开辟新的内存,而想指定结果到原来的y的内存,我们可以使用索引来进行替换操作。比如把x + y的结果通过[:]写进y对应的内存中。

    x = torch.tensor([1, 2])
    y = torch.tensor([3, 4])
    id_0 = id(y)
    y[:] = y + x
    print(id(y) == id_0) 
    # True
    

    另外,以下两种方式也可以索引到相同的内存:

    • torch.add(x, y, out=y)
    • y += x, y.add_(x)
    x = torch.tensor([1, 2])
    y = torch.tensor([3, 4])
    id_0 = id(y)
    torch.add(x, y, out=y) 
    # y += x, y.add_(x)
    print(id(y) == id_0) 
    # True
    

    clone() 与 detach() 对比

    Torch 为了提高速度,向量或是矩阵的赋值是指向同一内存的,这不同于 Matlab。如果需要保存旧的tensor即需要开辟新的存储地址而不是引用,可以用 clone() 进行深拷贝,

    首先我们来打印出来clone()操作后的数据类型定义变化:

    (1). 简单打印类型

    import torch
    
    a = torch.tensor(1.0, requires_grad=True)
    b = a.clone()
    c = a.detach()
    a.data *= 3
    b += 1
    
    print(a) # tensor(3., requires_grad=True)
    print(b)
    print(c)
    
    '''
    输出结果:
    tensor(3., requires_grad=True)
    tensor(2., grad_fn=<AddBackward0>)
    tensor(3.)  # detach()后的值随着a的变化出现变化
    '''
    

    grad_fn=<CloneBackward>,表示clone后的返回值是个中间变量,因此支持梯度的回溯。clone操作在一定程度上可以视为是一个identity-mapping函数。

    detach()操作后的tensor与原始tensor共享数据内存,当原始tensor在计算图中数值发生反向传播等更新之后,detach()的tensor值也发生了改变。

    注意: 在pytorch中我们不要直接使用id是否相等来判断tensor是否共享内存,这只是充分条件,因为也许底层共享数据内存,但是仍然是新的tensor,比如detach(),如果我们直接打印id会出现以下情况。

    import torch as t
    a = t.tensor([1.0,2.0], requires_grad=True)
    b = a.detach()
    #c[:] = a.detach()
    print(id(a))
    print(id(b))
    #140568935450520
    140570337203616
    

    显然直接打印出来的id不等,我们可以通过简单的赋值后观察数据变化进行判断。

    (2). clone()的梯度回传

    detach()函数可以返回一个完全相同的tensor,与旧的tensor共享内存,脱离计算图,不会牵扯梯度计算。

    而clone充当中间变量,会将梯度传给源张量进行叠加,但是本身不保存其grad,即值为None

    import torch
    a = torch.tensor(1.0, requires_grad=True)
    a_ = a.clone()
    y = a**2
    z = a ** 2+a_ * 3
    y.backward()
    print(a.grad) # 2
    z.backward()
    print(a_.grad)   # None. 中间variable,无grad
    print(a.grad) 
    '''
    输出:
    tensor(2.) 
    None
    tensor(7.) # 2*2+3=7
    '''
    

    使用torch.clone()获得的新tensor和原来的数据不再共享内存,但仍保留在计算图中,clone操作在不共享数据内存的同时支持梯度梯度传递与叠加,所以常用在神经网络中某个单元需要重复使用的场景下。

    通常如果原tensor的requires_grad=True,则:

    • clone()操作后的tensor requires_grad=True
    • detach()操作后的tensor requires_grad=False。
    import torch
    torch.manual_seed(0)
    
    x= torch.tensor([1., 2.], requires_grad=True)
    clone_x = x.clone() 
    detach_x = x.detach()
    clone_detach_x = x.clone().detach() 
    
    f = torch.nn.Linear(2, 1)
    y = f(x)
    y.backward()
    
    print(x.grad)
    print(clone_x.requires_grad)
    print(clone_x.grad)
    print(detach_x.requires_grad)
    print(clone_detach_x.requires_grad)
    '''
    输出结果如下:
    tensor([-0.0053, 0.3793])
    True
    None
    False
    False
    '''
    

    另一个比较特殊的是当源张量的 require_grad=False,clone后的张量 require_grad=True,此时不存在张量回传现象,可以得到clone后的张量求导。

    如下:

    import torch
    a = torch.tensor(1.0)
    a_ = a.clone()
    a_.requires_grad_() #require_grad=True
    y = a_ ** 2
    y.backward()
    print(a.grad) # None
    print(a_.grad) 
    '''
    输出:
    None
    tensor(2.)
    '''
    

    总结:

    torch.detach() —新的tensor会脱离计算图,不会牵扯梯度计算

    torch.clone() — 新的tensor充当中间变量,会保留在计算图中,参与梯度计算(回传叠加),但是一般不会保留自身梯度。

    原地操作(in-place, such as resize_ / resize_as_ / set_ / transpose_) 在上面两者中执行都会引发错误或者警告。

    引用官方文档的话:如果你使用了in-place operation而没有报错的话,那么你可以确定你的梯度计算是正确的。另外尽量避免in-place的使用。

    jsjbwy
    下一篇:没有了