| 46 | return ret |
| 47 | |
| 48 | def gumbel_softmax(logits: torch.Tensor, tau: float = 1, dim: int = -2) -> torch.Tensor: |
| 49 | gumbel_dist = torch.distributions.gumbel.Gumbel( |
| 50 | torch.tensor(0., device=logits.device, dtype=logits.dtype), |
| 51 | torch.tensor(1., device=logits.device, dtype=logits.dtype)) |
| 52 | gumbels = gumbel_dist.sample(logits.shape) |
| 53 | |
| 54 | gumbels = (logits + gumbels) / tau |
| 55 | y_soft = gumbels.softmax(dim) |
| 56 | |
| 57 | index = y_soft.max(dim, keepdim=True)[1] |
| 58 | y_hard = torch.zeros_like(logits, memory_format=torch.legacy_contiguous_format).scatter_(dim, index, 1.0) |
| 59 | ret = y_hard - y_soft.detach() + y_soft |
| 60 | |
| 61 | return ret |
| 62 | |
| 63 | class Fusion(nn.Module): |
| 64 | def __init__(self, in_dim_1, in_dim_2, out_dim, bias=False) -> None: |