Method get_head_mask

bert/modeling_utils.py:232–251

# Prepare head mask if needed # 1.0 in head_mask indicate we keep the head attention_probs has shape bsz x n_heads x N x N Arguments: head_mask: torch.Tensor or None: has shape [num_heads] or [num_hidden_layers x num_heads] num_hidden_layers:

(self, head_mask: Tensor, num_hidden_layers: int, is_attention_chunked: bool = False)

        return extended_attention_mask

    def get_head_mask(self, head_mask: Tensor, num_hidden_layers: int, is_attention_chunked: bool = False) -> Tensor:
        """
        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        attention_probs has shape bsz x n_heads x N x N
        Arguments:
            head_mask: torch.Tensor or None: has shape [num_heads] or [num_hidden_layers x num_heads]
            num_hidden_layers: int
        Returns:
             Tensor of shape shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
             or list with [None] for each layer
        """
        if head_mask is not None:
            head_mask = self._convert_head_mask_to_5d(head_mask, num_hidden_layers)
            if is_attention_chunked is True:
                head_mask = head_mask.unsqueeze(-1)
        else:
            head_mask = [None] * num_hidden_layers

        return head_mask

    def _convert_head_mask_to_5d(self, head_mask, num_hidden_layers):
        """-> [num_hidden_layers x batch x num_heads x seq_length x seq_length]"""