def forward(self,
src: Tensor,
mask: Optional[Tensor] = None,
src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
r"""Pass the input through the encoder layers in turn.
Args:
src: the sequence to the encoder (required).
mask: the mask for the src sequence (optional).
src_key_padding_mask: the mask for the src keys per batch (optional).
Shape:
see the docs in Transformer class.
"""
output = src
convert_to_nested = False
first_layer = self.layers[0]
src_key_padding_mask_for_layers = src_key_padding_mask
if isinstance(first_layer, torch.nn.TransformerEncoderLayer):
if (not first_layer.norm_first and not first_layer.training
and first_layer.self_attn.batch_first
and first_layer.self_attn._qkv_same_embed_dim
and first_layer.activation_relu_or_gelu
and first_layer.norm1.eps == first_layer.norm2.eps
and src.dim() == 3 and self.enable_nested_tensor
and src_key_padding_mask is not None
and torch._nested_tensor_from_mask_left_aligned(
src, src_key_padding_mask.logical_not())):
if src_key_padding_mask is not None and not output.is_nested and mask is None:
tensor_args = (
src,
first_layer.self_attn.in_proj_weight,
first_layer.self_attn.in_proj_bias,
first_layer.self_attn.out_proj.weight,
first_layer.self_attn.out_proj.bias,
first_layer.norm1.weight,
first_layer.norm1.bias,
first_layer.norm2.weight,
first_layer.norm2.bias,
first_layer.linear1.weight,
first_layer.linear1.bias,
first_layer.linear2.weight,
first_layer.linear2.bias,
)
if not torch.overrides.has_torch_function(tensor_args):
if not torch.is_grad_enabled() or all(
[not x.requires_grad for x in tensor_args]):
if output.is_cuda or 'cpu' in str(output.device):
convert_to_nested = True
output = torch._nested_tensor_from_mask(
output, src_key_padding_mask.logical_not())
src_key_padding_mask_for_layers = None
for mod in self.layers:
output = mod(output,
src_mask=mask,
src_key_padding_mask=src_key_padding_mask_for_layers)
if convert_to_nested:
output = output.to_padded_tensor(0.)
if self.norm is not None:
output = self.norm(output)
return output
def grad_test_func(inpt):
nt = torch._nested_tensor_from_mask(inpt, mask)
# This implicitly tests to_padded_tensor grads
return nt.to_padded_tensor(0)
def forward_scriptable(
self,
src_tokens,
src_lengths,
return_all_hiddens: bool = False,
):
"""
Args:
src_tokens (LongTensor): tokens in the source language of shape
`(batch, src_len)`
src_lengths (torch.LongTensor): lengths of each source sentence of
shape `(batch)`
return_all_hiddens (bool, optional): also return all of the
intermediate hidden states (default: False).
Returns:
dict:
- **encoder_out** (Tensor): the last encoder layer's output of
shape `(src_len, batch, embed_dim)`
- **encoder_padding_mask** (ByteTensor): the positions of
padding elements of shape `(batch, src_len)`
- **encoder_embedding** (Tensor): the (scaled) embedding lookup
of shape `(batch, src_len, embed_dim)`
- **encoder_states** (List[Tensor]): all intermediate
hidden states of shape `(src_len, batch, embed_dim)`.
Only populated if *return_all_hiddens* is True.
"""
if self.pre_encoder is not None:
x, src_lengths, encoder_padding_mask = self.pre_encoder(
src_tokens, src_lengths
)
else:
x, encoder_padding_mask = (
src_tokens,
~speech_utils.sequence_mask(src_lengths, src_tokens.size(1)),
)
has_pads = src_tokens.device.type == "xla" or encoder_padding_mask.any()
if self.fc0 is not None:
x = self.dropout_module(x)
x = self.fc0(x)
x = self.embed_scale * x
if self.embed_positions is not None:
# 0s in `~encoder_padding_mask` are used as pad_idx for positional embeddings
x = x + self.embed_positions((~encoder_padding_mask).int())
if self.layernorm_embedding is not None:
x = self.layernorm_embedding(x)
x = self.dropout_module(x)
if self.quant_noise is not None:
x = self.quant_noise(x)
# account for padding while computing the representation
if has_pads:
x = x * (1 - encoder_padding_mask.unsqueeze(-1).type_as(x))
# B x T x C -> T x B x C
x = x.transpose(0, 1)
encoder_states = []
fc_results = []
if return_all_hiddens:
encoder_states.append(x)
attn_mask = self.get_attn_mask(src_lengths)
# nested tensor and BT enable
layer = self.layers[0]
BT_flag = False
NT_flag = False
# torch version check, BT>=1.12.0 and NT>=1.13.0.dev20220613
# internal format is '1.13.0a0+fb'
# external format is '1.13.0.dev20220613'(cpu&gpu) for nightly or "1.11.0"(cpu) or '1.11.0+cu102'(gpu) for stable
BT_version = False
NT_version = False
if "fb" in torch.__version__:
BT_version = True
NT_version = True
else:
if "+" in torch.__version__:
torch_version = torch.__version__.split("+")[0]
else:
torch_version = torch.__version__
torch_version = torch_version.split(".")
int_version = (
int(torch_version[0]) * 1000
+ int(torch_version[1]) * 10
+ int(torch_version[2])
)
if len(torch_version) == 3:
if int_version >= 1120:
BT_version = True
if int_version >= 1131:
NT_version = True
elif len(torch_version) == 4:
if int_version >= 1130:
BT_version = True
# Consider _nested_tensor_from_mask_left_aligned is landed after "20220613"
if int_version >= 1131 or (
int_version == 1130 and torch_version[3][3:] >= "20220613"
):
NT_version = True
if (
BT_version
and x.dim() == 3
and layer.load_to_BT
and not layer.return_fc
and layer.can_use_fastpath
and not layer.training
and not layer.ever_training
and not layer.cfg_checkpoint_activations
):
# Batch first can not be justified but needs user to make sure
x = x.transpose(0, 1)
# Check mask conditions for nested tensor
if NT_version:
if (
encoder_padding_mask is not None
and torch._nested_tensor_from_mask_left_aligned(
x, encoder_padding_mask.logical_not()
)
):
if not torch.is_grad_enabled() or not x.requires_grad:
x = torch._nested_tensor_from_mask(
x, encoder_padding_mask.logical_not()
)
NT_flag = True
BT_flag = True
# encoder layers
if NT_flag:
processing_mask = None
else:
processing_mask = encoder_padding_mask
encoder_padding_mask_out = processing_mask if has_pads else None
for layer in self.layers:
lr = layer(
x,
encoder_padding_mask=encoder_padding_mask_out,
attn_mask=attn_mask,
)
if isinstance(lr, tuple) and len(lr) == 2:
x, fc_result = lr
else:
x = lr
fc_result = None
if return_all_hiddens and not torch.jit.is_scripting():
assert encoder_states is not None
encoder_states.append(x)
fc_results.append(fc_result)
# change back to non-nested and Batch second
if NT_flag:
x = x.to_padded_tensor(0.0)
if NT_flag or BT_flag:
x = x.transpose(0, 1)
if self.layer_norm is not None:
x = self.layer_norm(x)
# The Pytorch Mobile lite interpreter does not supports returning NamedTuple in
# `forward` so we use a dictionary instead.
# TorchScript does not support mixed values so the values are all lists.
# The empty list is equivalent to None.
return {
"encoder_out": [x], # T x B x C
"encoder_padding_mask": [encoder_padding_mask]
if encoder_padding_mask.any()
else [], # B x T
"encoder_embedding": [],
"encoder_states": encoder_states, # List[T x B x C]
"fc_results": fc_results, # List[T x B x C]
"src_tokens": [],
"src_lengths": [src_lengths], # List[B]
}
def _create_nested_tensor_from_mask(self, requires_grad=False):
data = torch.randn(2, 3, 4, requires_grad=requires_grad)
mask = torch.ones_like(data[:, :, 0]).bool()
return torch._nested_tensor_from_mask(data, mask)
def forward(self,
src: Tensor,
mask: Optional[Tensor] = None,
src_key_padding_mask: Optional[Tensor] = None) -> Tensor:
r"""Pass the input through the encoder layers in turn.
Args:
src: the sequence to the encoder (required).
mask: the mask for the src sequence (optional).
src_key_padding_mask: the mask for the src keys per batch (optional).
Shape:
see the docs in Transformer class.
"""
output = src
convert_to_nested = False
first_layer = self.layers[0]
src_key_padding_mask_for_layers = src_key_padding_mask
why_not_sparsity_fast_path = ''
str_first_layer = "self.layers[0]"
if not isinstance(first_layer, torch.nn.TransformerEncoderLayer):
why_not_sparsity_fast_path = f"{str_first_layer} was not TransformerEncoderLayer"
elif first_layer.norm_first:
why_not_sparsity_fast_path = f"{str_first_layer}.norm_first was True"
elif first_layer.training:
why_not_sparsity_fast_path = f"{str_first_layer} was in training mode"
elif not first_layer.self_attn.batch_first:
why_not_sparsity_fast_path = f" {str_first_layer}.self_attn.batch_first was not True"
elif not first_layer.self_attn._qkv_same_embed_dim:
why_not_sparsity_fast_path = f"{str_first_layer}.self_attn._qkv_same_embed_dim was not True"
elif not first_layer.activation_relu_or_gelu:
why_not_sparsity_fast_path = f" {str_first_layer}.activation_relu_or_gelu was not True"
elif not (first_layer.norm1.eps == first_layer.norm2.eps):
why_not_sparsity_fast_path = f"{str_first_layer}.norm1.eps was not equal to {str_first_layer}.norm2.eps"
elif not src.dim() == 3:
why_not_sparsity_fast_path = f"input not batched; expected src.dim() of 3 but got {src.dim()}"
elif not self.enable_nested_tensor:
why_not_sparsity_fast_path = "enable_nested_tensor was not True"
elif src_key_padding_mask is None:
why_not_sparsity_fast_path = "src_key_padding_mask was None"
elif (((not hasattr(self, "mask_check")) or self.mask_check)
and not torch._nested_tensor_from_mask_left_aligned(
src, src_key_padding_mask.logical_not())):
why_not_sparsity_fast_path = "mask_check enabled, and src and src_key_padding_mask was not left aligned"
elif output.is_nested:
why_not_sparsity_fast_path = "NestedTensor input is not supported"
elif mask is not None:
why_not_sparsity_fast_path = "src_key_padding_mask and mask were both supplied"
elif first_layer.self_attn.num_heads % 2 == 1:
why_not_sparsity_fast_path = "num_head is odd"
if not why_not_sparsity_fast_path:
tensor_args = (
src,
first_layer.self_attn.in_proj_weight,
first_layer.self_attn.in_proj_bias,
first_layer.self_attn.out_proj.weight,
first_layer.self_attn.out_proj.bias,
first_layer.norm1.weight,
first_layer.norm1.bias,
first_layer.norm2.weight,
first_layer.norm2.bias,
first_layer.linear1.weight,
first_layer.linear1.bias,
first_layer.linear2.weight,
first_layer.linear2.bias,
)
if torch.overrides.has_torch_function(tensor_args):
why_not_sparsity_fast_path = "some Tensor argument has_torch_function"
elif not (src.is_cuda or 'cpu' in str(src.device)):
why_not_sparsity_fast_path = "src is neither CUDA nor CPU"
elif torch.is_grad_enabled() and any(x.requires_grad
for x in tensor_args):
why_not_sparsity_fast_path = (
"grad is enabled and at least one of query or the "
"input/output projection weights or biases requires_grad")
if (not why_not_sparsity_fast_path) and (src_key_padding_mask
is not None):
convert_to_nested = True
# simplify on or after on 8/16/2022 to unconditionally call with mask_check=False
# we have established that either (1) the mask is OK with the check above,
# or (2) that we don't need a mask check with mask_check=False in the init
if not torch.jit.is_scripting():
output = torch._nested_tensor_from_mask(
output,
src_key_padding_mask.logical_not(),
mask_check=False)
else:
# When scripting, make a simpler call until the FC bar passes on 8/16/2022
output = torch._nested_tensor_from_mask(
output, src_key_padding_mask.logical_not())
src_key_padding_mask_for_layers = None
for mod in self.layers:
output = mod(output,
src_mask=mask,
src_key_padding_mask=src_key_padding_mask_for_layers)
if convert_to_nested:
output = output.to_padded_tensor(0.)
if self.norm is not None:
output = self.norm(output)
return output