Masked Softmax in PyTorch

masked_softmax.py

import torch
import torch.nn as nn


class MaskedSoftmax(nn.Module):
    def __init__(self):
        super(MaskedSoftmax, self).__init__()
        self.softmax = nn.Softmax(1)

    def forward(self, x, mask=None):
        """
        Performs masked softmax, as simply masking post-softmax can be
        inaccurate
        :param x: [batch_size, num_items]
        :param mask: [batch_size, num_items]
        :return:
        """
        if mask is not None:
            mask = mask.float()
        if mask is not None:
            x_masked = x * mask + (1 - 1 / mask)
        else:
            x_masked = x
        x_max = x_masked.max(1)[0]
        x_exp = (x - x_max.unsqueeze(-1)).exp()
        if mask is not None:
            x_exp = x_exp * mask.float()
        return x_exp / x_exp.sum(1).unsqueeze(-1)

@jiangnanhugo yes it was a trick to create infinite values back in the day. Nowadays we have masked_fill, which is a better solution.

Hello, I don't understand why did ‘mask’ be multiplied again in line 26, what's the impact if not do that?

It is to ensure that the sum used for normalization excludes non-masked elements. This implementation is not ideal for the latest pytorch (v1.6). I’ll update the code when I am available.

…

On Wed, Sep 9, 2020 at 12:23 AM cabbagecabbage ***@***.***> wrote: ***@***.**** commented on this gist. ------------------------------ Hello, I don't understand why did ‘mask’ be multiplied again in line 26, what's the impact if don't do that? — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub <https://gist.github.com/94f3ede72d1651b087a561cf80b306ca#gistcomment-3445986>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAHFGRQJNWOSUJGSCUC2NFLSEZD6VANCNFSM4JPYLRBQ> .

Hello, I found this post via Google.

I made a smaller, functional version that works for any tensor shape:

import torch

def masked_softmax(x, mask, **kwargs):
    x_masked = x.clone()
    x_masked[mask == 0] = -float("inf")

    return torch.softmax(x_masked, **kwargs)

If anyone else ever needs it :)