Model#

GeneralistModel#

class ofasys.model.ofa.GeneralistModel(cfg: Optional[GeneralistModelConfig] = None)[source]#

Initializes internal Module state, shared by both nn.Module and ScriptModule.

extract_features(src_tokens, src_lengths, prev_output_tokens, **kwargs)[source]#

Similar to forward but only return features.

Returns

the decoder’s features of shape (batch, tgt_len, embed_dim)
a dictionary with any model-specific outputs

Return type

tuple

forward(slots: List[Slot], features_only: bool = False, full_context_alignment: bool = False, alignment_layer: Optional[int] = None, alignment_heads: Optional[int] = None, return_all_hiddens: bool = False, return_encoder_out: bool = False, return_hf_dict: bool = False, return_all_attention_weights: bool = False)[source]#

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

get_normalized_probs(net_output: Tuple[Tensor, Optional[Dict[str, List[Optional[Tensor]]]]], log_probs: bool, sample: Optional[Dict[str, Tensor]] = None)[source]#: Get normalized probabilities (or log probs) from a net’s output.

max_decoder_positions()[source]#: Maximum length supported by the decoder.

max_positions()[source]#: Maximum length supported by the model.

upgrade_state_dict_named(state_dict, name)[source]#

Upgrade old state dicts to work with newer code.

Parameters

state_dict (dict) – state dictionary to upgrade, in place
name (str) – the state dict key corresponding to the current module

GeneralistModelConfig#

class ofasys.model.ofa.GeneralistModelConfig(_name: Union[str, NoneType] = None, activation_fn: ofasys.configure.constants.Choices = 'relu', dropout: float = 0.1, attention_dropout: float = 0.0, activation_dropout: float = 0.0, adaptive_input: bool = False, encoder: ofasys.module.transformer_config.EncDecBaseConfig = EncDecBaseConfig(_name=None, embed_path=None, embed_dim=512, ffn_embed_dim=2048, layers=6, attention_heads=8, normalize_before=False, learned_pos=False, layerdrop=0, layers_to_keep=None), max_source_positions: int = 1024, decoder: ofasys.module.transformer_config.DecoderConfig = DecoderConfig(_name=None, embed_path=None, embed_dim=512, ffn_embed_dim=2048, layers=6, attention_heads=8, normalize_before=False, learned_pos=False, layerdrop=0, layers_to_keep=None, input_dim=512, output_dim=512), max_target_positions: int = 1024, share_decoder_input_output_embed: bool = False, share_all_embeddings: bool = False, no_token_positional_embeddings: bool = False, adaptive_softmax_cutoff: Union[List[int], NoneType] = None, adaptive_softmax_dropout: float = 0.0, adaptive_softmax_factor: float = 4, layernorm_embedding: bool = False, tie_adaptive_weights: bool = False, tie_adaptive_proj: bool = False, no_scale_embedding: bool = False, checkpoint_activations: bool = False, offload_activations: bool = False, no_cross_attention: bool = False, cross_self_attention: bool = False, quant_noise: ofasys.module.transformer_config.QuantNoiseConfig = QuantNoiseConfig(_name=None, pq=0.0, pq_block_size=8, scalar=0.0), min_params_to_wrap: int = 100000000, char_inputs: bool = False, relu_dropout: float = 0.0, base_layers: Union[int, NoneType] = 0, base_sublayers: Union[int, NoneType] = 1, base_shuffle: Union[int, NoneType] = 1, export: bool = False, no_decoder_final_norm: bool = False, arch: str = 'base', encode_drop_path_rate: float = 0.0, decode_drop_path_rate: float = 0.0, attn_scale_factor: float = 2, freeze_encoder: bool = False, freeze_encoder_embedding: bool = False, freeze_decoder_embedding: bool = False, add_type_embedding: bool = True, entangle_position_embedding: bool = False, sync_bn: bool = False, scale_attn: bool = True, scale_fc: bool = True, scale_heads: bool = True, scale_resids: bool = False, checkpoint_adaptor_activations: bool = False, use_fused: bool = False, use_self_attn_bias: bool = True, adaptor: ofasys.adaptor.general.OFAAdaptorConfig = OFAAdaptorConfig(_name=None, text=TextAdaptorConfig(_name=None, is_active=False, layernorm_embedding=True, layernorm_position=True, add_type_embedding=True, entangle_position_embedding=False, no_scale_embedding=True, scale_embedding_gradient=1.0, dropout=None, embed_dim=None, num_attention_heads=None, encoder_layers=None, decoder_layers=None, max_position=None, use_self_attn_bias=None, share_attn_bias=None, token_bucket_size=256, share_input_output_embed=True, output_embed_dim=512, output_dim=None, output_bias=False), image_resnet=ImageResnetAdaptorConfig(_name=None, is_active=False, layernorm_embedding=True, layernorm_position=True, add_type_embedding=True, entangle_position_embedding=False, no_scale_embedding=True, scale_embedding_gradient=1.0, dropout=None, embed_dim=None, num_attention_heads=None, encoder_layers=None, decoder_layers=None, max_position=None, use_self_attn_bias=None, share_attn_bias=None, resnet_type='resnet152', resnet_drop_path_rate=0.0, sync_bn=False, freeze_resnet=False, image_bucket_size=42, pretrained_ckpt_path=''), image_patch_embed=ImagePatchEmbedAdaptorConfig(_name=None, is_active=False, layernorm_embedding=True, layernorm_position=True, add_type_embedding=True, entangle_position_embedding=False, no_scale_embedding=True, scale_embedding_gradient=1.0, dropout=None, embed_dim=768, num_attention_heads=None, encoder_layers=None, decoder_layers=None, max_position=None, use_self_attn_bias=None, share_attn_bias=None, image_size_width=224, image_size_height=224, patch_size_width=14, patch_size_height=14, add_cls_token=True), image_vqgan=ImageVqganAdaptorConfig(_name=None, is_active=False, layernorm_embedding=True, layernorm_position=True, add_type_embedding=True, entangle_position_embedding=False, no_scale_embedding=True, scale_embedding_gradient=1.0, dropout=None, embed_dim=None, num_attention_heads=None, encoder_layers=None, decoder_layers=None, max_position=None, use_self_attn_bias=None, share_attn_bias=None, code_image_size=256, code_bucket_size=42, vqgan_factor=8, vqgan_model_path='oss://ofasys/tasks/image_gen/vqgan/last.ckpt', vqgan_config_path='oss://ofasys/tasks/image_gen/vqgan/model.yaml', use_encode=True, code_entry_prefix='code'), audio_fbank=AudioFbankAdaptorConfig(_name=None, is_active=False, layernorm_embedding=True, layernorm_position=True, add_type_embedding=True, entangle_position_embedding=False, no_scale_embedding=True, scale_embedding_gradient=1.0, dropout=None, embed_dim=None, num_attention_heads=None, encoder_layers=None, decoder_layers=None, max_position=None, use_self_attn_bias=None, share_attn_bias=None, output_frame_dim=80, n_frames_per_step=1, is_transformer_layers=False, encoder_config=EncDecBaseConfig(_name=None, embed_path=None, embed_dim=512, ffn_embed_dim=2048, layers=6, attention_heads=8, normalize_before=False, learned_pos=False, layerdrop=0, layers_to_keep=None), encode_drop_path_rate=0.0, checkpoint_activations=False, min_params_to_wrap=100000000, attn_scale_factor=2, scale_attn=True, scale_fc=True, scale_heads=True, scale_resids=False, use_fused=True, prenet_layers=2, prenet_dim=256, prenet_dropout=0.5, postnet_conv_dim=512, postnet_conv_kernel_size=5, postnet_layers=5, postnet_dropout=0.5, use_mask=False, mask_length=10, mask_prob=0.65, mask_selection='static', mask_other=0, no_mask_overlap=False, mask_min_space=1, mask_channel_length=10, mask_channel_prob=0.0, mask_channel_selection='static', mask_channel_other=0, no_mask_channel_overlap=False, mask_channel_min_space=1, mask_channel_before=False, require_same_masks=True, mask_dropout=0.0), audio_tgt_fbank=AudioTargetFbankAdaptorConfig(_name=None, is_active=False, layernorm_embedding=True, layernorm_position=True, add_type_embedding=True, entangle_position_embedding=False, no_scale_embedding=True, scale_embedding_gradient=1.0, dropout=None, embed_dim=None, num_attention_heads=None, encoder_layers=None, decoder_layers=None, max_position=None, use_self_attn_bias=None, share_attn_bias=None, output_frame_dim=80, n_frames_per_step=1, conv_kernel_size=5, prenet_layers=2, prenet_dim=256, prenet_dropout=0.5, postnet_conv_dim=512, postnet_conv_kernel_size=5, postnet_layers=5, postnet_dropout=0.5), image_vit=ImageVitAdaptorConfig(_name=None, is_active=False, layernorm_embedding=True, layernorm_position=True, add_type_embedding=True, entangle_position_embedding=False, no_scale_embedding=True, scale_embedding_gradient=1.0, dropout=None, embed_dim=None, num_attention_heads=None, encoder_layers=None, decoder_layers=None, max_position=None, use_self_attn_bias=None, share_attn_bias=None, vit_type='vit_base', vit_drop_path_rate=0.0, image_bucket_size=42, pretrained_ckpt_path=''), motion_6d=Motion6dAdaptorConfig(_name=None, is_active=False, layernorm_embedding=True, layernorm_position=True, add_type_embedding=True, entangle_position_embedding=False, no_scale_embedding=True, scale_embedding_gradient=1.0, dropout=None, embed_dim=None, num_attention_heads=None, encoder_layers=None, decoder_layers=None, max_position=None, use_self_attn_bias=None, share_attn_bias=None, token_bucket_size=256, share_input_output_embed=True, output_embed_dim=512, output_dim=None, output_bias=False, max_data_dim=512, max_noise_levels=1024), video_image_sequence=VideoImageSequenceAdaptorConfig(_name=None, is_active=False, layernorm_embedding=True, layernorm_position=True, add_type_embedding=True, entangle_position_embedding=False, no_scale_embedding=True, scale_embedding_gradient=1.0, dropout=None, embed_dim=None, num_attention_heads=None, encoder_layers=None, decoder_layers=None, max_position=None, use_self_attn_bias=None, share_attn_bias=None, token_bucket_size=256)), share_attn_bias: bool = False, modal_ffn: bool = False, extra_models: ofasys.model.ofa.ExtraModelsConfig = ExtraModelsConfig(_name=None, pooling={}))[source]#

add_type_embedding: bool = True#

arch: str = 'base'#

attn_scale_factor: float = 2#

checkpoint_adaptor_activations: bool = False#

decode_drop_path_rate: float = 0.0#

encode_drop_path_rate: float = 0.0#

entangle_position_embedding: bool = False#

extra_models: ExtraModelsConfig = ExtraModelsConfig(_name=None, pooling={})#

freeze_decoder_embedding: bool = False#

freeze_encoder: bool = False#

freeze_encoder_embedding: bool = False#

modal_ffn: bool = False#

scale_attn: bool = True#

scale_fc: bool = True#

scale_heads: bool = True#

scale_resids: bool = False#

share_attn_bias: bool = False#

sync_bn: bool = False#

use_fused: bool = False#

use_self_attn_bias: bool = True#

TransformerEncoder#

class ofasys.model.transformer.TransformerEncoder(cfg, dictionary: Dictionary)[source]#

Transformer encoder consisting of cfg.encoder_layers layers. Each layer is a TransformerEncoderLayer.

Parameters

cfg (GeneralistModelConfig) – parsed command-line arguments
dictionary (Dictionary) – global dictionary

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(slots: List[Slot], return_all_hiddens: bool = False, return_all_attention_weights: bool = False)[source]#

Parameters

slots (List[Slot]) – preprocessed data
return_all_hiddens (bool, optional) – also return all of the intermediate hidden states (default: False).
return_all_attention_weights (bool, optional) – also return all attention weights (default: False).

Returns

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.
position_embeddings (Tensor): the position embedding lookup of shape (batch, src_len, embed_dim) - encoder_attention_weights (Tensor): attention weights of encoder’s self attention of shape (num_heads, batch_size, src_len, src_len). Only return if return_all_attention_weights and return_encoder_out are both True.

Return type

dict

max_positions()[source]#: Maximum input length supported by the encoder.

reorder_encoder_out(encoder_out: Dict[str, List[Tensor]], new_order)[source]#

Reorder encoder output according to new_order.

Parameters

encoder_out – output from the forward() method
new_order (LongTensor) – desired order

Returns

encoder_out rearranged according to new_order

TransformerDecoder#

class ofasys.model.transformer.TransformerDecoder(cfg, dictionary, no_encoder_attn=False)[source]#

Transformer decoder consisting of cfg.decoder_layers layers. Each layer is a TransformerDecoderLayer.

Parameters

cfg (GeneralistModelConfig) – arguments
dictionary (Dictionary) – decoding dictionary
no_encoder_attn (bool, optional) – whether to attend to encoder outputs (default: False).

Initializes internal Module state, shared by both nn.Module and ScriptModule.

extract_features(slots: List[Slot], encoder_out: Optional[Dict[str, List[Tensor]]], incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None, full_context_alignment: bool = False, alignment_layer: Optional[int] = None, alignment_heads: Optional[int] = None, return_all_hiddens: bool = False, return_all_attention_weights: bool = False)[source]#

Similar to forward but only return features.

Includes several features from “Jointly Learning to Align and Translate with Transformer Models” (Garg et al., EMNLP 2019).

Parameters

slots (List[Slot]) – preprocessed data.
encoder_out (optional, Dict[str, List[Tensor]]) – output from the encoder, used for encoder-side attention.
incremental_state (dict) – dictionary used for storing state during Incremental decoding.
full_context_alignment (bool, optional) – don’t apply auto-regressive mask to self-attention (default: False).
alignment_layer (int, optional) – return mean alignment over heads at this layer (default: last layer).
alignment_heads (int, optional) – only average alignment over this many heads (default: all heads).
return_all_hiddens (bool, optional) – also return all of the intermediate hidden states (default: False).
return_all_attention_weights (bool, optional) – also return all attention weights (default: False).

Returns

the decoder’s features of shape (batch, tgt_len, embed_dim).
a dictionary with decoder extra outputs.
- attn (List[Tensor]) : return specific attention weights.
- inner_states (List[Tensor]): all intermediate encoder hidden states of shape (tgt_len, batch, embed_dim).
- decoder_attentions (List[Tensor]): attention weights of decoder’s self attention of shape (num_heads, batch_size, tgt_len, tgt_len). Only return if return_all_attention_weights is True.
- cross_attentions (List[Tensor]): attention weights of decoder’s self attention of shape (num_heads, batch_size, src_len, tgt_len). Only return if return_all_attention_weights is True.

Return type

tuple

forward(slots: List[Slot], encoder_out: Optional[Dict[str, List[Tensor]]] = None, incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None, features_only: bool = False, full_context_alignment: bool = False, alignment_layer: Optional[int] = None, alignment_heads: Optional[int] = None, return_all_hiddens: bool = False, return_all_attention_weights: bool = False)[source]#

Parameters

slots (List[Slot]) – preprocessed data
encoder_out (optional, Dict[str, List[Tensor]]) – output from the encoder, used for encoder-side attention.
incremental_state (dict) – dictionary used for storing state during Incremental decoding
features_only (bool, optional) – only return features without applying output layer (default: False).
full_context_alignment (bool, optional) – don’t apply auto-regressive mask to self-attention (default: False).
alignment_layer (int, optional) – return mean alignment over heads at this layer (default: last layer).
alignment_heads (int, optional) – only average alignment over this many heads (default: all heads).
return_all_hiddens (bool, optional) – also return all of the intermediate hidden states (default: False).
return_all_attention_weights (bool, optional) – also return all attention weights (default: False).

Returns

the decoder’s output: the decoder’s features of shape (batch, tgt_len, embed_dim) if features_only is True, else return outputs from adaptor.
a dictionary with decoder extra outputs.
- attn (List[Tensor]) : return specific attention weights
- inner_states (List[Tensor]): all intermediate encoder hidden states of shape (tgt_len, batch, embed_dim),
- decoder_attentions (List[Tensor]): attention weights of decoder’s self attention of shape (num_heads, batch_size, tgt_len, tgt_len). Only return if return_all_attention_weights is True.
- cross_attentions (List[Tensor]): attention weights of decoder’s self attention of shape (num_heads, batch_size, src_len, tgt_len). Only return if return_all_attention_weights is True.

Return type

tuple

get_cross_pos_info(embed, tgt_pos_embed, src_pos_embed)[source]#: Compute abs position bias for cross attention.

max_positions()[source]#: Maximum output length supported by the decoder.

TransformerEncoderLayer#

class ofasys.module.transformer_layer.TransformerEncoderLayer(args, drop_path_rate=0.0)[source]#

Encoder layer block.

In the original paper each operation (multi-head attention or FFN) is postprocessed with: dropout -> add residual -> layernorm. In the tensor2tensor code they suggest that learning is more robust when preprocessing each layer with layernorm and postprocessing with: dropout -> add residual. We default to the approach in the paper, but the tensor2tensor approach can be enabled by setting cfg.encoder.normalize_before to True.

Parameters: args (argparse.Namespace) – parsed command-line arguments

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x, encoder_padding_mask: Optional[Tensor], attn_mask: Optional[Tensor] = None, self_attn_bias: Optional[Tensor] = None, need_attn: bool = False, modal_mask=None)[source]#

Parameters

x (Tensor) – input to the layer of shape (seq_len, batch, embed_dim)
encoder_padding_mask (ByteTensor) – binary ByteTensor of shape (batch, seq_len) where padding elements are indicated by 1.
attn_mask (ByteTensor) – binary tensor of shape (tgt_len, src_len), where tgt_len is the length of output and src_len is the length of input, though here both are equal to seq_len. attn_mask[tgt_i, src_j] = 1 means that when calculating the embedding for tgt_i, we exclude (mask out) src_j. This is useful for strided self-attention.
self_attn_bias (Tensor) –
need_attn (bool) –

Returns

encoded output of shape (seq_len, batch, embed_dim)

upgrade_state_dict_named(state_dict, name)[source]#: Rename layer norm states from …layer_norms.0.weight to …self_attn_layer_norm.weight and …layer_norms.1.weight to …final_layer_norm.weight

TransformerDecoderLayer#

class ofasys.module.transformer_layer.TransformerDecoderLayer(args, no_encoder_attn=False, add_bias_kv=False, add_zero_attn=False, drop_path_rate=0.0)[source]#

Decoder layer block.

In the original paper each operation (multi-head attention, encoder attention or FFN) is postprocessed with: dropout -> add residual -> layernorm. In the tensor2tensor code they suggest that learning is more robust when preprocessing each layer with layernorm and postprocessing with: dropout -> add residual. We default to the approach in the paper, but the tensor2tensor approach can be enabled by setting cfg.decoder_normalize_before to True.

Parameters

args (argparse.Namespace) – parsed command-line arguments
no_encoder_attn (bool, optional) – whether to attend to encoder outputs (default: False).
add_bias_kv (bool, optional) – (default: False).
add_zero_attn (bool, optional) – (default: False).
drop_path_rate (float, optional) – (default: 0.0).

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x, encoder_out: Optional[Tensor] = None, encoder_padding_mask: Optional[Tensor] = None, incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None, prev_self_attn_state: Optional[List[Tensor]] = None, prev_attn_state: Optional[List[Tensor]] = None, self_attn_mask: Optional[Tensor] = None, self_attn_padding_mask: Optional[Tensor] = None, need_attn: bool = False, need_head_weights: bool = False, self_attn_bias: Optional[Tensor] = None, cross_attn_bias: Optional[Tensor] = None, modal_mask=None)[source]#

Parameters

x (Tensor) – input to the layer of shape (seq_len, batch, embed_dim)
encoder_out –
encoder_padding_mask (ByteTensor, optional) – binary ByteTensor of shape (batch, src_len) where padding elements are indicated by 1.
incremental_state –
self_attn_mask –
self_attn_padding_mask –
need_attn (bool, optional) – return attention weights
need_head_weights (bool, optional) – return attention weights for each head (default: return average over heads).
self_attn_bias (Tensor, optional) – attention bias for self attention.
cross_attn_bias (Tensor, optional) – attenion bias for cross attention.

Returns

encoded output of shape (seq_len, batch, embed_dim)

upgrade_state_dict_named(state_dict, name)[source]#: Rename layer norm states from …layer_norms.0.weight to …self_attn_layer_norm.weight and …layer_norms.1.weight to …final_layer_norm.weight