174dd18905f12bfe0b418a4bad3575e6331478a4,python/dgl/nn/pytorch/glob.py,MultiHeadAttention,forward,#MultiHeadAttention#Any#Any#Any#Any#,663

Before Change

        e = e / np.sqrt(self.d_head)

        // generate mask
        mask = th.zeros(batch_size, max_len_x, max_len_mem).to(e.device)
        for i in range(batch_size):
            mask[i, :lengths_x[i], :lengths_mem[i]].fill_(1)
        mask = mask.unsqueeze(1)
        e.masked_fill_(mask == 0, -float("inf"))

        // apply softmax
        alpha = th.softmax(e, dim=-1)

After Change

        max_len_x = max(lengths_x)
        max_len_mem = max(lengths_mem)
        device = x.device
        lengths_x = th.tensor(lengths_x, dtype=th.int64, device=device)
        lengths_mem = th.tensor(lengths_mem, dtype=th.int64, device=device)

        queries = self.proj_q(x).view(-1, self.num_heads, self.d_head)
        keys = self.proj_k(mem).view(-1, self.num_heads, self.d_head)
        values = self.proj_v(mem).view(-1, self.num_heads, self.d_head)

        // padding to (B, max_len_x/mem, num_heads, d_head)
        queries = F.pad_packed_tensor(queries, lengths_x, 0)
        keys = F.pad_packed_tensor(keys, lengths_mem, 0)
        values = F.pad_packed_tensor(values, lengths_mem, 0)

        // attention score with shape (B, num_heads, max_len_x, max_len_mem)
        e = th.einsum("bxhd,byhd->bhxy", queries, keys)
        // normalize
        e = e / np.sqrt(self.d_head)

        // generate mask
        mask = _gen_mask(lengths_x, lengths_mem, max_len_x, max_len_mem)
        e = e.masked_fill(mask == 0, -float("inf"))

        // apply softmax
        alpha = th.softmax(e, dim=-1)
        // the following line addresses the NaN issue, see
        // https://github.com/dmlc/dgl/issues/2657
        alpha = alpha.masked_fill(mask == 0, 0.)

        // sum of value weighted by alpha
        out = th.einsum("bhxy,byhd->bxhd", alpha, values)
        // project to output

Instances