pytorch Embedding max_norm使用注意事项

pytorch Embedding

在官方文档中，Embedding的max_norm参数给出了一个注意事项：When max_norm is not None, Embedding’s forward method will modify the weight tensor in-place. Since tensors needed for gradient computations cannot be modified in-place, performing a differentiable operation on Embedding.weight before calling Embedding’s forward method requires cloning Embedding.weight when max_norm is not None。
大意是当Embedding层的max_norm参数不为None时，调用forward函数会原地修改Embedding的weight的值，所以如果要在调用forward函数之前对weight进行可微操作需要对weight进行复制。

n, d, m = 3, 5, 7
embedding = nn.Embedding(n, d, max_norm=True)
W = torch.randn((m, d), requires_grad=True)
idx = torch.tensor([1, 2])
a = embedding.weight.clone() @ W.t()  # weight must be cloned for this to be differentiable
b = embedding(idx) @ W.t()  # modifies weight in-place
out = (a.unsqueeze(0) + b.unsqueeze(1))
loss = out.sigmoid().prod()
loss.backward()

Embedding的forward会对weight进行修改的是因为在forward函数中调用的是F.embedding函数，F.embedding如果设置了max_norm的话，会对传入的weight进行原地的修改。在上述代码中，b = embedding(idx)调用了forward函数，a = embedding.weight @ W.t()在forward之前对embedding.weight使用了可微操作，所以a处需要复制，因为如果不复制，在反向传播的时候，经过b的运算，weight参数被修改了，之前记录的和反向传播的时候得到的参数不一致，会产生错误。如果交换ab两行的位置，也就是在forward之后再对embedding.weight调用可谓操作，就不需要复制了