about the loss in the second stage of training

[jun000697]

I followed readme to reproduce the project and when I do stage 2 training, the loss contains nan.

so,What should I do to resume normal training?

[jun000697]

What does the warning in this image mean? Does it mean that I didn't use a custom UNet (including blockfe) and instead trained using the officially defined UNet, or can this warning be ignored?

[indu1ge]

Regarding the abnormal loss issue, could you please provide the first-stage model you trained?
As for the warning about the U-Net weights loading, you can ignore it. Since the first-stage model does not contain the blockfe module, those weights will be newly initialized.

[jun000697]

I did not undergo the first stage of training and went directly to the second stage, so I do not have the weights from the first stage. 双控研2501班郭君君 ***@***.***

[indu1ge]

Hi,

Regarding the nan loss issue, you can try commenting out Lines 325 and 326 in train_s2.py.

Additionally, we have fixed several other issues. Please refer to my recent commit: 63999f6 (click the hash to view the changes).

After updating the code, the training should be normal, and the loss should be around 0.006.

Hope this helps!

[jun000697]

Thank you for your reply. I also noticed that the depth relationship in your images seems to be represented differently in colors. Why is the depth near blue and the depth far red? Also, my training speed seems a bit slow. My device is an RTX 4090 graphics card. Could you tell me how to improve the training speed? 双控研2501班郭君君 ***@***.***

[indu1ge]

Hi,

We choose different colormaps to better highlight near details. You can change the colormap in the visualization code.

For the slow training speed, maybe you could try :

• Enabling mixed precision
• Change a device with larger VRAM so that you can increase batch size
• Reduce the number of training iterations.

[jun000697]

Hi, even after I commented out those two lines of code as you suggested, the loss = Nan still occurred. 双控研2501班郭君君 ***@***.***

[indu1ge]

It's suggested to undergo the first-stage training and use the first-stage model to initialize the model weights.

[jun000697]

Why? 双控研2501班郭君君 ***@***.***

[indu1ge]

The current setting and training code is suited for our two-stage training with the first stage learning basic structure and the second stage refining details. If you directly train the second-stage model, maybe you should modify training settings accordingly.

[jun000697]

Hello, when I ran the first stage of training, I also encountered a loss of NaN. How can I resolve this? 双控研2501班郭君君 ***@***.***

[indu1ge]

We have tested the training code, and the loss output is normal. Maybe you could check the data, weights and environment setup.