gumbel softmax
-
用于处理argmax不可导的情况
-
解决思路:引入gumbel分布。在前向传播中使用argmax,后向梯度回传中使用gumbel_softmax计算
-
代码
def gumbel_softmax(logits: Tensor, tau: float = 1, hard: bool = False, eps: float = 1e-10, dim: int = -1) -> Tensor:
。。。
gumbels = (
-torch.empty_like(logits, memory_format=torch.legacy_contiguous_format).exponential_().log()
) # ~Gumbel(0,1)
gumbels = (logits + gumbels) / tau # ~Gumbel(logits,tau)
y_soft = gumbels.softmax(dim)
if hard:
# Straight through.
index = y_soft.max(dim, keepdim=True)[1]
y_hard = torch.zeros_like(logits, memory_format=torch.legacy_contiguous_format).scatter_(dim, index, 1.0)
ret = y_hard - y_soft.detach() + y_soft
else:
# Reparametrization trick.
ret = y_soft
return ret
gumbel_softmax中引入了温度t, t越小,softmax就越接近One-hot。为了训练稳定性,一般t会取一个比较大的数字,然后逐步缩小。
内容转载自gumbel softmax