search

BAAI-DCAI / SegVol

The official code for "SegVol: Universal and Interactive Volumetric Medical Image Segmentation".
MIT License
245 stars 22 forks source link

UnboundLocalError: local variable 'loss' referenced before assignment #16

Open TousenKaname opened 8 months ago

TousenKaname commented 8 months ago

Thanks for your excellent work! But I encounter a bug when I set the 'use_pseudo_label=False'.

Traceback (most recent call last):
  File "/nvme/gawang/miniconda3/envs/Segvol/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap
    fn(i, *args)
  File "/nvme/gawang/git_repo/SegVol/train.py", line 193, in main_worker
    epoch_loss, iter_num = train_epoch(args, segvol_model, train_dataloader, optimizer, scheduler, epoch, rank, gpu, iter_num)
  File "/nvme/gawang/git_repo/SegVol/train.py", line 91, in train_epoch
    loss_step_avg += loss.item()
UnboundLocalError: local variable 'loss' referenced before assignment

The loss is not defined...

TousenKaname commented 8 months ago

And I found that 'args.use_pseudo_label' only affect the loss calculate, it means we also have to dounsupervised_forward method, consuming computing waste...

Yuxin-Du-Lab commented 8 months ago

I'm sorry for that. The args.use_pseudo_label is still unfinished. The original released code only supports training with pseudo labels.

But you can disable the pseudo-label training by discarding all the code related to the pseudo-label in the training process. SegVol will perform better training with only ground truth labels in specific categories.