首先给出模型结构图：然后给出完整的代码：importloggingfromdataclassesimportdataclass,fieldfromtypingimportDict,List,Optional,Tupleimportnumpyasnpimporttorchimporttorch.nnasnnfromomegaconfimportIIfromfairseqimportutilsfromfairseq.data.data_utilsimportcompute_mask_indicesfromfairseq.data.dictionaryimportDictionaryfromfairseq.dataclassimportChoiceEnum,FairseqDataclassfromfairseq.modelsimportBaseFairseqModel,register_modelfromfairseq.models.wav2vec.wav2vec2import(EXTRACTOR_MODE_CHOICES,MASKING_DISTRIBUTION_CHOICES,LAYER_TYPE_CHOICES,ConvFeatureExtractionModel,TransformerEncoder,)fromfairseq.modulesimportGradMultiply,LayerNormfromfairseq.tasks.hubert_pretrainingimport(HubertPretrainingConfig,HubertPretrainingTask,)logger=logging.getLogger(__name__)@dataclassclassHubertConfig(FairseqDataclass):label_rate:float=II("task.label_rate")extractor_mode:EXTRACTOR_MODE_CHOICES=field(default="default",metadata={"help":"mode for feature extractor. default has a single group ""norm with d groups in the first conv block, whereas layer_norm ""has layer norms in every block (meant to use with normalize=True)"},)encoder_layers:int=field(default=12,metadata={"help":"num encoder layers in the transformer"})encoder_embed_dim:int=field(default=768,metadata={"help":"encoder embedding dimension"})encoder_ffn_embed_dim:int=field(default=3072,metadata={"help":"encoder embedding dimension for FFN"})encoder_attention_heads:int=field(default=12,metadata={"help":"num encoder attention heads"})activation_fn:ChoiceEnum(utils.get_available_activation_fns())=field(default="gelu",metadata={"help":"activation function to use"})layer_type:LAYER_TYPE_CHOICES=field(default="transformer",metadata={"help":"layer type in encoder"})# dropoutsdropout:float=field(default=0.1,metadata={"help":"dropout probability for the transformer"},)attention_dropout:float=field(default=0.1,metadata={"help":"dropout probability for attention weights"},)activation_dropout:float=field(default=0.0,metadata={"help":"dropout probability after activation in FFN"},)encoder_layerdrop:float=field(default=0.0,metadata={"help":"probability of dropping a tarnsformer layer"},)dropout_input:float=field(default=0.0,metadata={"help":"dropout to apply to the input (after feat extr)"},)dropout_features:float=field(default=0.0,metadata={"help":"dropout to apply to the features (after feat extr)"},)final_dim:int=field(default=0,metadata={"help":"project final representations and targets to this many ""dimensions. set to encoder_embed_dim is = 0"},)untie_final_proj:bool=field(default=False,metadata={"help":"use separate projection for each target"},)layer_norm_first:bool=field(default=False,metadata={"help":"apply layernorm first in the transformer"},)conv_feature_layers:str=field(default="[(512,10,5)] + [(512,3,2)] * 4 + [(512,2,2)] * 2",metadata={"help":"string describing convolutional feature extraction ""layers in form of a python list that contains ""[(dim, kernel_size, stride), ...]"},)conv_bias:bool=field(default=False,metadata={"help":"include bias in conv encoder"})logit_temp:float=field(default=0.1,metadata={"help":"temperature to divide logits by"})target_glu:bool=field(default=False,metadata={"help":"adds projection + glu to targets"})feature_grad_mult:float=field(default=1.0,metadata={"help":"multiply feature extractor var grads by this"},)# maskingmask_length:int=field(default=10,metadata={"help":"mask length"})mask_prob:float=field(default=0.65,metadata={"help":"probability of replacing a token with mask"},)mask_selection:MASKING_DISTRIBUTION_CHOICES=field(default="static",metadata={"help":"how to choose mask length"})mask_other:float=field(default=0,metadata={"help":"secondary mask argument ""(used for more complex distributions), ""see help in compute_mask_indicesh"},)no_mask_overlap:bool=field(default=False,metadata={"help":"whether to allow masks to overlap"})mask_min_space:int=field(default=1,metadata={"help":"min space between spans (if no overlap is enabled)"},)# channel maskingmask_channel_length:int=field(default=10,metadata={"help":"length of the mask for features (channels)"},)mask_channel_prob:float=field(default=0.0,metadata={"help":"probability of replacing a feature with 0"},)mask_channel_selection:MASKING_DISTRIBUTION_CHOICES=field(default="static",metadata={"help":"how to choose mask length for channel masking"},)mask_channel_other:float=field(default=0,metadata={"help":"secondary mask argument ""(used for more complex distributions), ""see help in compute_mask_indicesh"},)no_mask_channel_overlap:bool=field(default=False,metadata={"help":"whether to allow channel masks to overlap"},)mask_channel_min_space
