PyTorch: loss.backward() keeps running for days

hxt18258480551: <p>I’m using the pretrained model from deepspeech_pytorch (<a href="https://github.com/SeanNaren/deepspeech.pytorch" class="inline-onebox" rel="noopener nofollow ugc">GitHub - SeanNaren/deepspeech.pytorch: Speech Recognition using DeepSpeech2.</a>) to do some speech recognition tasks.</p>
<p>However, the <code>loss.backward()</code> keeps running for days in training phase.</p>
<p>The code block of training phase:</p>
<pre data-code-wrap="Python"><code class="lang-plaintext">adversary_model = AdversaryModel( ... )

adv_trainset = WavDataset(adv_trainset_wav_path_list, model.labels, adv_train_spect_parser)
adv_trainset_dataloader = AudioAdvTrainDataLoader(dataset=adv_trainset, batch_size=FLAGS.batch_size, pin_memory=True, shuffle=True)
model.train()

# Because my task is something like adversary example, 
# all the params I need to train is only the "delta", 
# which is declared in addition as nn.Parameter(), 
# but not the params inner the neural network. 
opt = torch.optim.Adam([adversary_model.delta], lr=FLAGS.adv_lr)

adversary_model.to(FLAGS.device)
adversary_model.model.to(FLAGS.device)

for ep in range(FLAGS.adv_train_epochs):
    print(f'\nTrain epoch = { ep }\n', flush=True)
    adversary_model.train()
    for (batch_idx, batch) in enumerate(adv_trainset_dataloader):
        opt.zero_grad()
        

        # A lot of codes to input data batch into the model and compute loss ...

        loss *= adversary_model.factor_loss
        
        print(f'Adv train epoch = { ep }, batch_idx = { batch_idx }, loss.backward() start ...', flush=True)
        loss.backward()
        print(f'Adv train epoch = { ep }, batch_idx = { batch_idx }, loss.backward() finish ...', flush=True)
        
        print(f'Adv train epoch = { ep }, batch_idx = { batch_idx }, opt.step() start ...', flush=True)
        opt.step()
        print(f'Adv train epoch = { ep }, batch_idx = { batch_idx }, opt.step() finish ...', flush=True)
</code></pre>
<p>And the stdout:</p>
<pre data-code-wrap="shell"><code class="lang-plaintext">Train epoch = 0

Adv train epoch = 0, batch_idx = 0, loss.backward() start ...
Adv train epoch = 0, batch_idx = 0, loss.backward() finish ...
Adv train epoch = 0, batch_idx = 0, opt.step() start ...
Adv train epoch = 0, batch_idx = 0, opt.step() finish ...
Adv train epoch = 0, batch_idx = 1, loss.backward() start ...
</code></pre>
<p>Without any warnings or errors, it just keep running for days and never stop …</p>
<p>My current <code>batch_size</code> is 32, And I’ve tried to reduce the <code>batch_size</code> to 16, 8, 4, 2, 1. Then I found out the issue occurred around the 40th sample.</p>
<p>And I found out that it seems like the <code>convolution reverb</code> function makes this problem happen, because when I remove this step, everything works fine:</p>
<pre data-code-wrap="Python"><code class="lang-plaintext">def convolution_reverb(speech : Tensor, rir : Tensor):
    speech = speech.unsqueeze(0)
    rir = rir.unsqueeze(0)
    
    rir = rir / torch.linalg.vector_norm(rir, ord=2)
    rir = torch.flip(rir, [1])
    
    speech = torch.nn.functional.pad(speech, (rir.shape[1]-1, 0))
    reverbed = torch.nn.functional.conv1d(speech[None, ...], rir[None, ...])[0]

    return reverbed[0]
</code></pre>
<p>This <code>convolution_reverb</code> function is called inner the <code>parse_audio</code> phase, which is in the <code>__getitem__</code> function in my dataset class. This function is to simulate room reverberation for clean audio. The input arg <code>speech</code> is the input audio tensor (time domain), and the <code>rir</code> is the sample in some open-source RIR dataset (also time domain), which is preloaded to CUDA. The code is copied from docs of <code>torchaudio</code> (<a href="https://pytorch.org/audio/0.10.0/tutorials/audio_data_augmentation_tutorial.html#simulating-room-reverberation" class="inline-onebox" rel="noopener nofollow ugc">Audio Data Augmentation — Torchaudio 0.10.0 documentation</a>)</p>
<p>I’ve put all the tensor computation on GPU, and I’ve checked the OS memory (64G in total) and GPU memory (48G in total, 2 GeForce RTX 3090) are both enough. I’ve tried to debug this code but got nothing useful.</p>
<p>And the version information of my codes and environments:</p>
<ul>
<li>OS: Linux 5.4.0 / Ubuntu 18.04</li>
<li>Python: 3.7.14</li>
<li>torch: 1.10.0+cu111</li>
<li>torchaudio: 0.10.0+cu111</li>
</ul>
<p>Could anyone help me with why this keeps running or any good ways to debug the backward() function?</p>
<p>Thank you so much !!</p>

ptrblck: <p>Could you update PyTorch and all sublibraries to the latest version and check if the code is still hanging, please?</p>

hxt18258480551: <p>Thank you for your advice! I’ll try it soon.</p>

hxt18258480551: <aside class="quote no-group" data-username="hxt18258480551" data-post="1" data-topic="191776">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/hxt18258480551/48/65436_2.png" class="avatar"> hxt18258480551:</div>
<blockquote>
<p>nings or errors, it just keep running for days and never s</p>
</blockquote>
</aside>
<p>After update everything, the issue is gone, thank you very much !</p>
<p>New versions:</p>
<ul>
<li>Python 3.8.13</li>
<li>pytorch-lightning 2.1.1</li>
<li>torch 2.0.1</li>
<li>torchaudio 2.0.2</li>
</ul>

ptrblck: <p>Cool! NIT: the latest PyTorch version is <code>2.1.0</code> in case you want to update.</p>