Is it safe to reuse the same PyTorch loss module across multiple environments in a single training iteration?

autograd
1

I’m training a model where each training iteration processes two independent mini-batches (call them environment A and environment B). Both environments contribute to the same optimization step. I’m currently reusing the same loss modules ( Pytorch’s MSE and MONAI’s DiceCELoss) for both environments :

import torch
import torch.nn as nn
from monai.losses import DiceCELoss

model = MyModel()

optimizer = torch.optim.Adam(model.parameters())

mse_loss = nn.MSELoss()

dice_loss = DiceCELoss(
    include_background=False,
    sigmoid=True,
    softmax=False,
)

for batch_A, batch_B in zip(loader_A, loader_B):

    optimizer.zero_grad()

    #  Environment A 
    seg_logits_A, recon_A = model(batch_A["image"])

    seg_loss_A = dice_loss(
        seg_logits_A,
        batch_A["mask"],
    )

    recon_loss_A = mse_loss(
        recon_A,
        batch_A["image"],
    )

    # Environment B 
    seg_logits_B, recon_B = model(batch_B["image"])

    seg_loss_B = dice_loss(
        seg_logits_B,
        batch_B["mask"],
    )

    recon_loss_B = mse_loss(
        recon_B,
        batch_B["image"],
    )

    total_loss = (
        seg_loss_A +
        seg_loss_B +
        recon_loss_A +
        recon_loss_B
    )

    total_loss.backward()
    optimizer.step()

Both nn.MSELoss and DiceCELoss are instantiated only once and reused for both environments.

My questions are:

  1. Does autograd correctly accumulate gradients from both calls to the same loss instance?

  2. Is there any reason to instead instantiate separate losses ?

My understanding is that these loss modules are stateless, so reusing the same instance should be equivalent to creating multiple identical instances, but I’d like to confirm that there are no subtle autograd or gradient accumulation issues with this pattern. Thank you in advance.

2

Hi!

  1. Does autograd correctly accumulate gradients from both calls to the same loss instance?

Yes.

  1. Is there any reason to instead instantiate separate losses ?

No.

My understanding is that these loss modules are stateless, so reusing the same instance should be equivalent to creating multiple identical instances

Correct.

I don’t think there’s any problem with what you’re doing. It would be redundant to have different loss modules for loader_A and loader_B. As you mentioned, the only case where it would make a difference is if those loss modules have a state (trainable parameters or buffers like in BatchNorm), and you want this state to remain environment-specific. But since your loss modules are stateless you’re good.

3

Thank you very much!