torch.nn.CTCLoss using with multi-gpu

Viewed 16

I did CRNN with pytorch using https://github.com/Lornatang/CRNN-PyTorch

That codes can execute with single GPU, no problem.

But when I modified codes that they can use 2 GPUs(with adding nn.DataParallel()) they get same error like this:

RuntimeError: input_lengths must be of size batch_size

It always occured with traceback about nn.CTCLoss() method like below:

Traceback (most recent call last):
  File "train.py", line 370, in <module>
    main()
  File "train.py", line 84, in main
    train(model, train_dataloader, criterion, optimizer, epoch, scaler, writer)
  File "train.py", line 227, in train
    loss = criterion(output_log_probs, target, images_lengths, target_length)
  File "/home/gpudev/miniconda3/envs/gpu38/lib/python3.8/site-packages/torch/nn/parallel/parallel_apply.py", line 61, in _worker
    output = module(*input, **kwargs)
  File "/home/gpudev/miniconda3/envs/gpu38/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/gpudev/miniconda3/envs/gpu38/lib/python3.8/site-packages/torch/nn/modules/loss.py", line 1743, in forward
    return F.ctc_loss(log_probs, targets, input_lengths, target_lengths, self.blank, self.reduction,
  File "/home/gpudev/miniconda3/envs/gpu38/lib/python3.8/site-packages/torch/nn/functional.py", line 2599, in ctc_loss
    return torch.ctc_loss(
RuntimeError: input_lengths must be of size batch_size

I wonder torch.nn.CTCLoss cannot use with 2 or more GPUs

Thanks

0 Answers
Related