torch中的inplace操作问题解决方法

原创已于 2022-10-31 17:22:21 修改 · 1.3k 阅读

4 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #pytorch #机器学习

于 2022-10-31 16:01:46 首次发布

深度学习遇到的问题专栏收录该内容

1 篇文章

订阅专栏

本文介绍了PyTorch中的inplace操作，它是一种高效内存使用方法，但使用不当会在计算梯度或反传时出错。列举了典型的inplace操作，如+=、*=和不合适的原地切片赋值。通过出错案例对比，还给出判断inplace操作的方法，即看切片[]位置。

提示：文章写完后，目录可以自动生成，如何生成可参考右边的帮助文档

一、inplace操作是什么？

可以参考：PyTorch中的In-place操作是什么？为什么要避免使用这种操作？
简单来说，就是你写的代码采用了一种更为高效的内存使用方法，有点像copy(),deepcopy(),切片等操作的关系，如果你不小心使用了inplace操作的话，则在torch的计算梯度autograd()或者反传backward()位置处出现以下提示：
Runtime Error: Found an in-place operation that changes one the variables needed by gradient computation

典型的inplace操作有哪些？

1.+=，*=等操作
例如：

 a[:,0,:,:]+=b[:,0,:,:]

2.不合适的原地切片并赋值
例如：

 a[:,0,:,:]=do_sth(a[:,0,:,:])#do_sth是一个函数

这也是我这次出错的原因
这也是我这次出错的原因

二、出错案例

代码如下（示例）：

def forward(self, x):
        x_ = torch.zeros(self.layer(x[..., 0]).shape + (steps,), device=x.device)
        for step in range(steps):
            x_[..., step] = self.layer(x[..., step])
        if self.bn is not None:
            for i in range(steps):
                x_[:,:,:,:,i] = self.bn(x_[:,:,:,:,i])
        return x_

正确的代码如下：

def forward(self, x):

        x_ = torch.zeros(self.layer(x[..., 0]).shape + (steps,), device=x.device)
        for step in range(steps):
            x_[..., step] = self.layer(x[..., step])

        x_2 = torch.zeros(self.layer(x[..., 0]).shape + (steps,), device=x.device)
        if self.bn is not None:
            for i in range(steps):
                x_2[:,:,:,:,i] = self.bn(x_[:,:,:,:,i])
        return x_2 if self.bn is not None else x_