modeling_utils.py中temperature_sampling采样的问题

[MrSworder]

这是我运行run_pretrain.py时的运行报错，显示在temperature_sampling函数中pred_ids = probs.cpu().multinomial(probs.size()[1],replacement=False)这一句存在问题，由于multinomial输入的weights只能是1维或2维，而这里的probs的shape是【32，128，21128】所以出错。请问这里应该如何修改？

python run_pretraining.py --data_dir=dataset/ --vocab_path=prev_trained_model/electra_tiny/vocab.txt --data_name=electra --config_path=prev_trained_model/electra_tiny/config.json --output_dir=outputs/

03/31/2022 11:38:38 - INFO - root - samples_per_epoch: 173385 03/31/2022 11:38:38 - INFO - root - device: cuda , distributed training: False, 16-bits training: False 03/31/2022 11:38:38 - INFO - model.tokenization_utils - Model name 'prev_trained_model/electra_tiny/vocab.txt' not found in model shortcut name list (). Assuming 'prev_trained_model/electra_tiny/vocab.txt' is a path or url to a directory containing tokenizer files. 03/31/2022 11:38:38 - INFO - model.tokenization_utils - Didn't find file prev_trained_model/electra_tiny/added_tokens.json. We won't load it. 03/31/2022 11:38:38 - INFO - model.tokenization_utils - Didn't find file prev_trained_model/electra_tiny/special_tokens_map.json. We won't load it. 03/31/2022 11:38:38 - INFO - model.tokenization_utils - Didn't find file prev_trained_model/electra_tiny/tokenizer_config.json. We won't load it. 03/31/2022 11:38:38 - INFO - model.tokenization_utils - loading file prev_trained_model/electra_tiny/vocab.txt 03/31/2022 11:38:38 - INFO - model.tokenization_utils - loading file None 03/31/2022 11:38:38 - INFO - model.tokenization_utils - loading file None 03/31/2022 11:38:38 - INFO - model.tokenization_utils - loading file None 03/31/2022 11:38:38 - INFO - model.configuration_utils - loading configuration file prev_trained_model/electra_tiny/config.json 03/31/2022 11:38:38 - INFO - model.configuration_utils - Model config { "attention_probs_dropout_prob": 0.1, "directionality": "bidi", "disc_weight": 50, "embedding_size": 312, "finetuning_task": null, "gen_weight": 1.0, "generator_size": "1/4", "hidden_act": "gelu", "hidden_dropout_prob": 0.1, "hidden_size": 312, "initializer_range": 0.02, "intermediate_size": 1200, "layer_norm_eps": 1e-12, "max_position_embeddings": 512, "num_attention_heads": 12, "num_hidden_layers": 4, "num_labels": 2, "output_attentions": false, "output_hidden_states": false, "pooler_fc_size": 312, "pooler_num_attention_heads": 12, "pooler_num_fc_layers": 3, "pooler_size_per_head": 128, "pooler_type": "first_token_transform", "pruned_heads": {}, "temperature": 1.0, "torchscript": false, "type_vocab_size": 2, "vocab_size": 21128 }

03/31/2022 11:38:41 - INFO - root - ***** Running training ***** 03/31/2022 11:38:41 - INFO - root - Num examples = 693540 03/31/2022 11:38:41 - INFO - root - Batch size = 32 03/31/2022 11:38:41 - INFO - root - Num steps = 21673 03/31/2022 11:38:41 - INFO - root - warmup_steps = 2167 03/31/2022 11:38:41 - INFO - root - Loading training examples for dataset/corpus/train/electra_file_0.json 03/31/2022 11:38:43 - INFO - root - Loading complete! torch.Size([32, 128, 21128]) Traceback (most recent call last): File "run_pretraining.py", line 363, in main() File "run_pretraining.py", line 276, in main outputs = model(input_ids=input_ids, File "/home/zorro/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl return forward_call(*input, **kwargs) File "/home/zorro/NLP/中文预训练语言模型/electra_pytorch_chinese-master/model/modeling_electra.py", line 735, in forward sample_ids = temperature_sampling(g_logits,self.config.temperature) File "/home/zorro/NLP/中文预训练语言模型/electra_pytorch_chinese-master/model/modeling_utils.py", line 763, in temperature_sampling pred_ids = probs.cpu().multinomial(probs.size()[1],replacement=False) RuntimeError: prob_dist must be 1 or 2 dim