TypeError: forward() got an unexpected keyword argument 'src_lengths'

[qiyea]

Hi @peikexin9

When I pretrain the model by using the script ./command/pretrain/pretrain_10k.sh, I got an error.

2022-03-08 16:53:39 | INFO | fairseq.trainer | begin training epoch 1
2022-03-08 16:53:39 | INFO | fairseq_cli.train | Start iterating over samples
Traceback (most recent call last):
  File "train.py", line 14, in <module>
    cli_main()
  File "/data/binVul/trex-main/fairseq_cli/train.py", line 496, in cli_main
    distributed_utils.call_main(cfg, main)
  File "/data/binVul/trex-main/fairseq/distributed/utils.py", line 369, in call_main
    main(cfg, **kwargs)
  File "/data/binVul/trex-main/fairseq_cli/train.py", line 173, in main
    valid_losses, should_stop = train(cfg, trainer, task, epoch_itr)
  File "/usr/local/miniconda3/envs/trex2/lib/python3.8/contextlib.py", line 75, in inner
    return func(*args, **kwds)
  File "/data/binVul/trex-main/fairseq_cli/train.py", line 284, in train
    log_output = trainer.train_step(samples)
  File "/usr/local/miniconda3/envs/trex2/lib/python3.8/contextlib.py", line 75, in inner
    return func(*args, **kwds)
  File "/data/binVul/trex-main/fairseq/trainer.py", line 669, in train_step
    loss, sample_size_i, logging_output = self.task.train_step(
  File "/data/binVul/trex-main/fairseq/tasks/fairseq_task.py", line 475, in train_step
    loss, sample_size, logging_output = criterion(model, sample)
  File "/usr/local/miniconda3/envs/trex2/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/data/binVul/trex-main/fairseq/criterions/trex.py", line 62, in forward
    output = model(**sample["net_input"], masked_code=masked_code, masked_value=masked_value)[0]
  File "/usr/local/miniconda3/envs/trex2/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
TypeError: forward() got an unexpected keyword argument 'src_lengths'

[RobinHan24]

the same error occurs as you. So how to solve it @peikexin9, thanks a lot.

[peikexin9]

Hi @RobinHan24 @qiyea Thanks again for your interests! One possible reason is recently I am actively migrating the code to be compatible with torchscript, so that it can be deployed without installing PyTorch. This might introduce many bugs. For example, I removed **kwargs in the forward function of fairseq/models/trex/model.py so that unused parameter like src_lengths are no longer available. I just updated the code by adding src_lengths back as one of the optional input parameters. See if that works?

[RobinHan24]

Thanks again. But there is another problem ''TypeError: forward() got an unexpected keyword argument 'masked_code' '' and '' 'TypeError: forward() got an unexpected keyword argument 'masked_value' ''. I tried to modify the code, but it didn't work. @peikexin9

[peikexin9]

@RobinHan24 Thanks for letting me know the issue. I am pretty sure my recent updates to make the code compatible to torchscript/jit compilable break it, since I have not tested pretraining after I make the updates. But I am too busy this week to fix it. I will work on it next week.

[RobinHan24]

Anyway, thank you for your sharing. But I haven't found the Bi-Lstm code which is used to deal with the dynamic value in your project as mentioned in this paper, could you please tell me where it is. And why there is 4 type（byte1 byte2 byte3 byte4） value about non-byte value? Thanks a lot. @peikexin9

[peikexin9]

Hi @RobinHan24 I just fixed the pretraining compatibility after migrating to torchscript. See if now it can run? Regarding your followup questions, I replaced the original LSTM with CNNs (

trex/fairseq/modules/trex_encoder.py

Line 65 in dd5384a

self.byte_combine = ByteCombineCNN(1, embed_dim)

) as I found it is faster and did not have performance drops. But if you want, you can use bi-LSTM by replacing with following code and rearrange your tensor shape in forward function a bit (just to the shape that the LSTM accepts):
nn.LSTM(self.embedding_dim, self.embedding_dim, 1, batch_first=True,bidirectional=True)
self.lstm_fc = nn.Linear(self.embedding_dim * 2, self.embedding_dim)
For byte1-4, as it is learning on most 32-bit arch (ARM32, MIPS32, x86-32), so I have this work-around to only use the least significant 4 bytes for x64 dataset to save more GPU space. Feel free to prepare your dataset with byte1-8.