pytorch 中是否有 theano.tensor.switch 的模拟？

Question

我想强制将矢量中低于特定阈值的所有元素归零。我想这样做，这样我仍然可以通过非零梯度传播梯度。

例如，在theano中我可以这样写：

B = theano.tensor.switch(A < .1, 0, A)

pytorch 中有解决方案吗？

Answer 1

我不认为 switch 在 PyTorch 中是默认实现的。但是，您可以通过 extending the torch.autograd.Function

在 PyTorch 中定义自己的函数

所以，switch 函数看起来像

class switchFunction(Function):
    @staticmethod
    def forward(ctx, flag, value, tensor):
        ctx.save_for_backward(flag)
        tensor[flag] = value
        return tensor

    @staticmethod
    def backward(ctx, grad_output):
        flag, = ctx.saved_variables
        grad_output[flag] = 0
        return grad_output
switch = switchFunction.apply

现在，您只需将 switch 称为 switch(A < 0.1, 0, A)

编辑

实际上有一个函数可以做到这一点。它被称为Threshold。你可以像

一样使用它

import torch.nn as nn
m = nn.Threshold(0.1, 0)
B = m(A)

Answer 2

从 pytorch 0.4+ 开始，您可以使用 torch.where 轻松完成（参见 doc,Merged PR）

它和 Theano 一样简单。举个例子看看自己：

import torch
from torch.autograd import Variable

x = Variable(torch.arange(0,4), requires_grad=True) # x     = [0 1 2 3]
zeros = Variable(torch.zeros(*x.shape))             # zeros = [0 0 0 0]

y = x**2                         # y = [0 1 4 9]
z = torch.where(y < 5, zeros, y) # z = [0 0 0 9]

# dz/dx = (dz/dy)(dy/dx) = (y < 5)(0) + (y ≥ 5)(2x) = 2x(x**2 ≥ 5) 
z.backward(torch.Tensor([1.0])) 
x.grad # (dz/dx) = [0 0 0 6]

pytorch 中是否有 theano.tensor.switch 的模拟？

Is there analog of theano.tensor.switch in pytorch?

theano

deep-learning

pytorch

编辑