Pytorch 中关于in-place operation的小总结

Pytorch 中关于in-place operation的小总结 问题引入

有些时候，我们需要对Tensor中的某些部分进行一些重新赋值的 *** 作（即in-place operation)，但是如果该步骤涉及到一些需要计算梯度的叶子变量 (即requires_grad=True)时，Pytorch会报如下错误：a leaf Variable that requires grad has been used in an in-place operation, 原因是这种in-place operation是一种原位 *** 作，是无法计算梯度的。比如下面这个简单的例子。

import torch
import torch.nn as nn

创建一个简单的 1d batch normalization layer

bn=nn.BatchNorm1d(10) # 10维的特征

创建一个简单的leaf tensor,同时要求计算梯度

a=torch.tensor(0.5, requires_grad=True)

检验是否为leaf

a.is_leaf
# 返回 True

进行in-place operation, 我希望将batch normalization层的权重参数设置为a.

bn.weight.fill_(a)
# 报错
# RuntimeError: a leaf Variable that requires grad is being used in an in-place operation.

解决方案

一个非常简单的方案就是，将这个 in-place operation放置在torch.no_grad()下面，这样的话，这一步骤不会参与梯度传播，自然也就不会报错。

with torch.no_grad():
	bn.weight.fill_(a)

我们来看一下赋值后的weight，没有出现问题。

Parameter containing:
tensor([0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000,
        0.5000], requires_grad=True)