gumbelmodule.py

import torch
import torch.nn.functional as F
from torch.autograd import Variable

"""
Gumbel Softmax Sampler
Requires 2D input [batchsize, number of categories]

Does not support sinlge binary category. Use two dimensions with softmax instead.
"""

class GumbleSoftmax(torch.nn.Module):
    def __init__(self, hard=False):
        super(GumbleSoftmax, self).__init__()
        self.hard = hard
        self.gpu = False
        
    def cuda(self):
        self.gpu = True
    
    def cpu(self):
        self.gpu = False
        
    def sample_gumbel(self, shape, eps=1e-10):
        """Sample from Gumbel(0, 1)"""
        noise = torch.rand(shape)
        noise.add_(eps).log_().neg_()
        noise.add_(eps).log_().neg_()
        if self.gpu:
            return Variable(noise).cuda()
        else:
            return Variable(noise)

    def sample_gumbel_like(self, template_tensor, eps=1e-10):
    #    uniform_samples_tensor = template_tensor.clone().uniform_()   # 返回一个和源张量同shape、dtype和device的张量，与源张量不共享数据内存，但提供梯度的回溯 从均匀分布中抽样数值进行填充
        uniform_samples_tensor = template_tensor.clone().normal_()
     #   print(uniform_samples_tensor)
        gumble_samples_tensor = -uniform_samples_tensor
    #    gumble_samples_tensor = - torch.log(eps - torch.log(uniform_samples_tensor + eps))
        return gumble_samples_tensor

    def gumbel_softmax_sample(self, logits, temperature):
        """ Draw a sample from the Gumbel-Softmax distribution"""
        dim = logits.size(-1)
        gumble_samples_tensor = self.sample_gumbel_like(logits.data)
        gumble_trick_log_prob_samples = logits + Variable(gumble_samples_tensor)
        soft_samples = F.softmax(gumble_trick_log_prob_samples / temperature, dim)
        return soft_samples
    
    def gumbel_softmax(self, logits, temperature, hard=False):
        """Sample from the Gumbel-Softmax distribution and optionally discretize.
        Args:
        logits: [batch_size, n_class] unnormalized log-probs
        temperature: non-negative scalar
        hard: if True, take argmax, but differentiate w.r.t. soft sample y
        Returns:
        [batch_size, n_class] sample from the Gumbel-Softmax distribution.
        If hard=True, then the returned sample will be one-hot, otherwise it will
        be a probabilitiy distribution that sums to 1 across classes
        """
        y = self.gumbel_softmax_sample(logits, temperature)
        
        if hard:
            _, max_value_indexes = y.data.max(1, keepdim=True)
            y_hard = logits.data.clone().zero_().scatter_(1, max_value_indexes, 1) #scatter_(input, dim, index, src)将src中数据根据index中的索引按照dim的方向填进input中

            y = Variable(y_hard - y.data) + y
        return y
        
    def forward(self, logits, temp=1, force_hard=False):
        samplesize = logits.size()
        
        if self.training and not force_hard:
            
            return self.gumbel_softmax(logits, temperature=1, hard=False)
        else:
            
            return self.gumbel_softmax(logits, temperature=1, hard=True)