UnboundLocalError: local variable 'l1_loss' referenced before assignment in Doctr's compute_loss function

[ep0p]

Bug description

While training a Doctr model with my own dataset, I encountered an UnboundLocalError in the compute_loss function of the differentiable_binarization module.

Code snippet to reproduce the bug

python references/detection/train_pytorch.py ../dataset_detection/train ../dataset_detection/val db_resnet50 --epochs 5 --device 0

Error traceback

Train set loaded in 63.19s (11871 samples in 5935 batches)
Training loss: 1.21584:   2%|█▊                                                                                  | 131/5935 [01:24<1:02:43,  1.54it/s]
Traceback (most recent call last):                                                                               | 131/5935 [01:24<1:02:04,  1.56it/s]
  File "/home/epop/dataset/doctr/references/detection/train_pytorch.py", line 481, in <module>
    main(args)
  File "/home/epop/dataset/doctr/references/detection/train_pytorch.py", line 388, in main
    fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, amp=args.amp)
  File "/home/epop/dataset/doctr/references/detection/train_pytorch.py", line 126, in fit_one_epoch
    train_loss = model(images, targets)["loss"]
  File "/home/epop/anaconda3/envs/ds3.9/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/epop/anaconda3/envs/ds3.9/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/epop/dataset/doctr/doctr/models/detection/differentiable_binarization/pytorch.py", line 216, in forward
    loss = self.compute_loss(logits, thresh_map, target)
  File "/home/epop/dataset/doctr/doctr/models/detection/differentiable_binarization/pytorch.py", line 286, in compute_loss
    return l1_loss + focal_scale * focal_loss + dice_loss
UnboundLocalError: local variable 'l1_loss' referenced before assignment

Environment

Collecting environment information...

DocTR version: 0.9.0a0
TensorFlow version: N/A
PyTorch version: 2.3.1+cu121 (torchvision 0.18.1+cu121)
OpenCV version: 4.10.0
OS: Ubuntu 22.04.4 LTS
Python version: 3.9.19
Is CUDA available (TensorFlow): N/A
Is CUDA available (PyTorch): Yes
CUDA runtime version: 11.5.119
GPU models and configuration:
GPU 0: NVIDIA GeForce RTX 3090
GPU 1: NVIDIA GeForce RTX 3090

Nvidia driver version: 555.42.02
cuDNN version: Could not collect

Deep Learning backend

>>> from doctr.file_utils import is_tf_available, is_torch_available
>>> print(f"is_tf_available: {is_tf_available()}")
is_tf_available: False
>>> print(f"is_torch_available: {is_torch_available()}")
is_torch_available: True

[felixdittrich92]

Hey @ep0p 👋,

This will only happen if we couldn't build the target mask/s
reasons for this:

• annotations (coordinates) are wrong
• a sample in the batch has no annotations
• all annotations of an sample was filtered because the are to small (less than 3x3 pixels)

[ep0p]

Hi @felixdittrich92,

I am in the process of checking my dataset but it will take a while, i'll come back to you as soon as i find the issue.

[felixdittrich92]

This is strange your dataset runs on my end (with the normal labels.json not normalized_labels.json)

(doctr-dev) felix@felix-Z790-AORUS-MASTER:~/Desktop/doctr$ USE_TORCH=1 python3 /home/felix/Desktop/doctr/references/detection/train_pytorch.py /home/felix/Desktop/X_TEST/ds_150pdfs/ds_150_det/train /home/felix/Desktop/X_TEST/ds_150pdfs/ds_150_det/val db_resnet50 --pretrained
Namespace(train_path='/home/felix/Desktop/X_TEST/ds_150pdfs/ds_150_det/train', val_path='/home/felix/Desktop/X_TEST/ds_150pdfs/ds_150_det/val', arch='db_resnet50', name=None, epochs=10, batch_size=2, device=None, save_interval_epoch=False, input_size=1024, lr=0.001, weight_decay=0, workers=None, resume=None, test_only=False, freeze_backbone=False, show_samples=False, wb=False, push_to_hub=False, pretrained=True, rotation=False, eval_straight=False, sched='poly', amp=False, find_lr=False, early_stop=False, early_stop_epochs=5, early_stop_delta=0.01)
Validation set loaded in 0.328s (217 samples in 109 batches)
Train set loaded in 0.9588s (864 samples in 432 batches)
Training loss: 0.19342: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 432/432 [02:15<00:00,  3.18it/s]
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 109/109 [00:35<00:00,  3.03it/s]
Validation loss decreased inf --> 0.159816: saving state...
Epoch 1/10 - Validation loss: 0.159816 (Recall: 79.07% | Precision: 83.29% | Mean IoU: 62.00%)

[ep0p]

Did it work for the full training process? For me it can complete several epochs and then it breaks. Like this:

Validation set loaded in 0.784s (217 samples in 109 batches)
Train set loaded in 2.013s (864 samples in 432 batches)
Training loss: 0.223406: 100%|███████████████████████████████████| 432/432 [03:53<00:00,  1.85it/s]
100%|████████████████████████████████████████████████████████████| 109/109 [01:06<00:00,  1.65it/s]
Validation loss decreased inf --> 0.201296: saving state...
Epoch 1/5 - Validation loss: 0.201296 (Recall: 54.04% | Precision: 60.53% | Mean IoU: 52.00%)
Training loss: 0.163602: 100%|███████████████████████████████████| 432/432 [03:49<00:00,  1.89it/s]
100%|████████████████████████████████████████████████████████████| 109/109 [01:06<00:00,  1.64it/s]
Validation loss decreased 0.201296 --> 0.164273: saving state...
Epoch 2/5 - Validation loss: 0.164273 (Recall: 66.53% | Precision: 72.86% | Mean IoU: 58.00%)
Training loss: 0.214939:  36%|████████████▍                      | 154/432 [01:20<02:25,  1.90it/s]
Traceback (most recent call last):████████▍                      | 154/432 [01:20<02:22,  1.95it/s]
  File "/home/epop/dataset_prep/doctr/references/detection/train_pytorch.py", line 481, in <module>
    main(args)
  File "/home/epop/dataset_prep/doctr/references/detection/train_pytorch.py", line 388, in main
    fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, amp=args.amp)
  File "/home/epop/dataset_prep/doctr/references/detection/train_pytorch.py", line 109, in fit_one_epoch
    for images, targets in pbar:
  File "/home/epop/anaconda3/envs/ds3.9/lib/python3.9/site-packages/tqdm/std.py", line 1181, in __iter__
    for obj in iterable:
  File "/home/epop/anaconda3/envs/ds3.9/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 631, in __next__
    data = self._next_data()
  File "/home/epop/anaconda3/envs/ds3.9/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 1346, in _next_data
    return self._process_data(data)
  File "/home/epop/anaconda3/envs/ds3.9/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 1372, in _process_data
    data.reraise()
  File "/home/epop/anaconda3/envs/ds3.9/lib/python3.9/site-packages/torch/_utils.py", line 705, in reraise
    raise exception
UnboundLocalError: Caught UnboundLocalError in DataLoader worker process 10.
Original Traceback (most recent call last):
  File "/home/epop/anaconda3/envs/ds3.9/lib/python3.9/site-packages/torch/utils/data/_utils/worker.py", line 308, in _worker_loop
    data = fetcher.fetch(index)  # type: ignore[possibly-undefined]
  File "/home/epop/anaconda3/envs/ds3.9/lib/python3.9/site-packages/torch/utils/data/_utils/fetch.py", line 51, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/home/epop/anaconda3/envs/ds3.9/lib/python3.9/site-packages/torch/utils/data/_utils/fetch.py", line 51, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/home/epop/dataset_prep/doctr/doctr/datasets/datasets/base.py", line 67, in __getitem__
    img_transformed, target[class_name] = self.sample_transforms(img, bboxes)
  File "/home/epop/dataset_prep/doctr/doctr/transforms/modules/base.py", line 56, in __call__
    x, target = t(x, target)
  File "/home/epop/anaconda3/envs/ds3.9/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/epop/anaconda3/envs/ds3.9/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/epop/dataset_prep/doctr/doctr/transforms/modules/pytorch.py", line 87, in forward
    target[:, [0, 2]] = offset[0] + target[:, [0, 2]] * raw_shape[-1] / img.shape[-1]
UnboundLocalError: local variable 'offset' referenced before assignment

[felixdittrich92]

Could you try to add here:

doctr/doctr/transforms/modules/pytorch.py

Line 80 in 1cea7d8

if target is not None:

if target is not None:
    target = np.clip(target, 0, 1)

Looks like a super stupid rounding issue which should normally not happen with the provided code this clips small overflows like 1.0000....1 to 1.0

[felixdittrich92]

(doctr-dev) felix@felix-Z790-AORUS-MASTER:~/Desktop/doctr$ USE_TORCH=1 python3 /home/felix/Desktop/doctr/references/detection/train_pytorch.py /home/felix/Desktop/X_TEST/ds_150pdfs/ds_150_det/train /home/felix/Desktop/X_TEST/ds_150pdfs/ds_150_det/val db_resnet50 --pretrained
Namespace(train_path='/home/felix/Desktop/X_TEST/ds_150pdfs/ds_150_det/train', val_path='/home/felix/Desktop/X_TEST/ds_150pdfs/ds_150_det/val', arch='db_resnet50', name=None, epochs=10, batch_size=2, device=None, save_interval_epoch=False, input_size=1024, lr=0.001, weight_decay=0, workers=None, resume=None, test_only=False, freeze_backbone=False, show_samples=False, wb=False, push_to_hub=False, pretrained=True, rotation=False, eval_straight=False, sched='poly', amp=False, find_lr=False, early_stop=False, early_stop_epochs=5, early_stop_delta=0.01)
Validation set loaded in 0.3329s (217 samples in 109 batches)
Train set loaded in 0.9803s (864 samples in 432 batches)
Training loss: 0.157946: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 432/432 [02:14<00:00,  3.21it/s]
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 109/109 [00:35<00:00,  3.10it/s]
Validation loss decreased inf --> 0.162874: saving state...
Epoch 1/10 - Validation loss: 0.162874 (Recall: 72.94% | Precision: 78.23% | Mean IoU: 60.00%)
Training loss: 0.123682: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 432/432 [02:14<00:00,  3.20it/s]
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 109/109 [00:35<00:00,  3.04it/s]
Validation loss decreased 0.162874 --> 0.142662: saving state...
Epoch 2/10 - Validation loss: 0.142662 (Recall: 83.28% | Precision: 86.91% | Mean IoU: 64.00%)
Training loss: 0.219536: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 432/432 [02:15<00:00,  3.19it/s]
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 109/109 [00:35<00:00,  3.09it/s]
Validation loss decreased 0.142662 --> 0.128305: saving state...
Epoch 3/10 - Validation loss: 0.128305 (Recall: 83.40% | Precision: 87.20% | Mean IoU: 65.00%)
Training loss: 0.175725: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 432/432 [02:16<00:00,  3.17it/s]
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 109/109 [00:35<00:00,  3.07it/s]
Validation loss decreased 0.128305 --> 0.119365: saving state...
Epoch 4/10 - Validation loss: 0.119365 (Recall: 87.68% | Precision: 90.51% | Mean IoU: 67.00%)
Training loss: 0.110322: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 432/432 [02:15<00:00,  3.18it/s]
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 109/109 [00:35<00:00,  3.06it/s]
Validation loss decreased 0.119365 --> 0.114303: saving state...
Epoch 5/10 - Validation loss: 0.114303 (Recall: 85.55% | Precision: 88.71% | Mean IoU: 65.00%)
Training loss: 0.124551: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 432/432 [02:17<00:00,  3.15it/s]
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 109/109 [00:36<00:00,  3.03it/s]
Validation loss decreased 0.114303 --> 0.111153: saving state...
Epoch 6/10 - Validation loss: 0.111153 (Recall: 84.54% | Precision: 88.14% | Mean IoU: 65.00%)
Training loss: 0.209552: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 432/432 [02:15<00:00,  3.18it/s]
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 109/109 [00:35<00:00,  3.10it/s]
Validation loss decreased 0.111153 --> 0.106859: saving state...
Epoch 7/10 - Validation loss: 0.106859 (Recall: 87.10% | Precision: 90.02% | Mean IoU: 67.00%)
Training loss: 0.107883: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 432/432 [02:13<00:00,  3.24it/s]
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 109/109 [00:35<00:00,  3.09it/s]
Validation loss decreased 0.106859 --> 0.103459: saving state...
Epoch 8/10 - Validation loss: 0.103459 (Recall: 86.66% | Precision: 89.70% | Mean IoU: 66.00%)
Training loss: 0.0908364: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 432/432 [02:14<00:00,  3.20it/s]
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 109/109 [00:35<00:00,  3.10it/s]
Validation loss decreased 0.103459 --> 0.101387: saving state...
Epoch 9/10 - Validation loss: 0.101387 (Recall: 88.01% | Precision: 90.85% | Mean IoU: 67.00%)
Training loss: 0.224623: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 432/432 [02:13<00:00,  3.23it/s]
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 109/109 [00:35<00:00,  3.10it/s]
Validation loss decreased 0.101387 --> 0.101059: saving state...
Epoch 10/10 - Validation loss: 0.101059 (Recall: 87.28% | Precision: 90.19% | Mean IoU: 67.00%)

[ep0p]

@felixdittrich92 all good for me as well with your suggested modif.
I will run some tests on larger datasets to see if it is not breaking but i suspect it will be fine.

Thanks a lot for your time and patience!

[aidevveloper]

Hi Guys!

.....................
File "/home/epop/dataset/doctr/doctr/models/detection/differentiable_binarization/pytorch.py", line 286, in compute_loss
return l1_loss + focal_scale * focal_loss + dice_loss
UnboundLocalError: local variable 'l1_loss' referenced before assignment
.....................

I solved the above error with the following change:

  def compute_loss(...):
        ........................
        # Initialize loss components to zero
        focal_loss = torch.tensor(0.0, device=out_map.device)
        dice_loss = torch.tensor(0.0, device=out_map.device)
        l1_loss = torch.tensor(0.0, device=out_map.device)
        focal_scale = 10.0

        if torch.any(seg_mask):
            # Focal loss
            bce_loss = F.binary_cross_entropy_with_logits(out_map, seg_target, reduction="none")
       .........................

[felixdittrich92]

Hi @Harry-KIT 👋

Careful here !
As you can see from the above discussion this error happens if something with your dataset isn't correct (empty target mask for example).
To check this you can save the target masks (should be used with a bs of 1):

prob_map = torch.sigmoid(out_map)
thresh_map = torch.sigmoid(thresh_map)

targets = self.build_target(target, out_map.shape[1:], False)  # type: ignore[arg-type]

seg_target, seg_mask = torch.from_numpy(targets[0]), torch.from_numpy(targets[1])
seg_target, seg_mask = seg_target.to(out_map.device), seg_mask.to(out_map.device)
thresh_target, thresh_mask = torch.from_numpy(targets[2]), torch.from_numpy(targets[3])
thresh_target, thresh_mask = thresh_target.to(out_map.device), thresh_mask.to(out_map.device)

import os
import cv2
os.makedirs("debug_masks", exist_ok=True)

def save_mask(tensor, name):
    if isinstance(tensor, torch.Tensor):
        mask = tensor.detach().cpu().numpy()
    else:
        mask = tensor

    if np.all(mask == 1) and name in ["seg_target", "thresh_mask"]:
        raise ValueError(f"[ERROR] Mask '{name}' is empty !\n\nTARGET:\n{target}")

    if mask.ndim == 4 and mask.shape[1] == 1:
        for i in range(mask.shape[0]):
            save_mask(mask[i, 0], f"{name}_{i}")
        return

    mask = np.clip(mask, 0, 1) * 255
    mask = mask.astype(np.uint8)

    cv2.imwrite(os.path.join("debug_masks", f"{name}.png"), mask)


save_mask(seg_mask, "seg_mask")
save_mask(seg_target, "seg_target")
save_mask(thresh_mask, "thresh_mask")
save_mask(thresh_target, "thresh_target")
save_mask(prob_map[0, 0], "prob_map")
save_mask(thresh_map[0, 0], "thresh_map")

Init the losses with 0 is no solution and would lead the model to wrong predictions because if there is an empty target there is nothing to "compare" and you would signal "blind" that the models prediction is really good because of a much lower loss -> backward.

Hope this helps 👍

Best,
Felix

[aidevveloper]

Hi @felixdittrich92, thank you for your efforts. I appreciate it.

I used CORD dataset, and I faced that error. But it works when I changed as i mentioned above.That is why I though it would be a solution. I never say I am surely true. I completely agree with your opinion.

..............
Epoch 6/50 - Validation loss: 0.165845 (Recall: 80.64% | Precision: 76.70% | Mean IoU: 58.00%)
Training loss: 0.215301 | LR: 0.00086: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 912/912 [02:52<00:00, 5.30it/s]
Epoch 7/50 - Training loss: 0.200824 | LR: 0.00086
Validation loss: 0.238143: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 241/241 [00:38<00:00, 6.30it/s]
Validation loss decreased 0.165845 --> 0.157576: saving state...
Epoch 7/50 - Validation loss: 0.157576 (Recall: 80.13% | Precision: 77.23% | Mean IoU: 60.00%)
Training loss: 0.273771 | LR: 0.000841009: 95%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████▊ | 866/912 [02:42<00:08, 5.31it/s]
Traceback (most recent call last):
File "/doctr/references/detection/train_pytorch.py", line 589, in
main(args)
File "/doctr/references/detection/train_pytorch.py", line 472, in main
train_loss, actual_lr = fit_one_epoch(
File "/doctr/references/detection/train_pytorch.py", line 131, in fit_one_epoch
train_loss = model(images, targets)["loss"]
File "/miniconda3/envs/OnnxTR/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/miniconda3/envs/OnnxTR/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
return forward_call(*args, **kwargs)
File "/doctr/doctr/models/detection/differentiable_binarization/pytorch.py", line 221, in forward
loss = self.compute_loss(logits, thresh_map, target)
File "/doctr/doctr/models/detection/differentiable_binarization/pytorch.py", line 289, in compute_loss
return l1_loss + focal_scale * focal_loss + dice_loss
UnboundLocalError: local variable 'l1_loss' referenced before assignment