Adjusting LandCoverSeg example for 2 instead of 5 classes

[jducnuigeen]

I've been trying to train a model on my own data using your example LandCoverSeg example with the aim to train a more pure segmentation model.
I've closely matched my custom input dataset to match the folder structure and image types for the images/masks/train/val/tests.
My dataset only has two classes : background, and buildings. Your original code running on the Landseg data had 5 classes.
I've also made some changes to hopefully accommodate the code for my two classes:

1. In src->datamodules->datasets I modified landcoverseg_dataset.py
   changing
   mask = np.zeros((*label.shape[:2], 5), dtype=np.uint8)
   to
   mask = np.zeros((*label.shape[:2], 2), dtype=np.uint8)

2. In src->models I modified segmenter_visualisate_utils.py
   changing the palatte from
   palette = [
   0, 0, 0,
   255,0,0,
   0,255,0,
   0,0,255,
   255,255,0, ]
   to
   palette = [
   0, 0, 0,
   1, 0, 0, ]

3. and of course changed the config, with mainly just class names and paths.

It tries to run, but I gets to /LandCoverSeg/src/losses/focal_dice_loss.py
and gives this error:

File "/content/drive/MyDrive/ML5/LandCoverSeg/src/losses/focal_dice_loss.py", line 21, in forward
    loss = torch.mean(loss, dim=(2,3)) * self.weights
RuntimeError: The size of tensor a (2) must match the size of tensor b (5) at non-singleton dimension 1

and i see that there are 5 torch tensor weight hardcoded on line 15 of that file, but I don't understand where they came from, or how I modify them for just two classes?
self.weights = torch.tensor([[0.34523039, 23.35764161, 0.60372157, 3.09578992, 12.32148783]]).cuda()

Are there other places I need to adjust to get this to work with 2 classes?
Thoughts?

[sepherro]

The weights correspond to the percentage of pixels for each class, as some classes will be underrepresented - for example, there will be (in general) far fewer pixels corresponding to roads than woods in a typical image. The weights make sure those pesky neural networks don't take the shortcut and actually make an effort to account for these underrepresented classes in the final segmentation result. Not sure what your classes are, but you can change it to a vector of two even numbers for starters, and it should work:

self.weights = torch.tensor([[0.5, 0.5]]).cuda()

[jducnuigeen]

Thank you for everyone's help. I managed to get this to work with your help - even a model trained on only 2200 images. I'll now work on writing up a tutorial for how I did this with my own imagery and share here somewhere in the discussion section.

[jducnuigeen]

As you can see above I got this workflow to successfully train the first time.
I now want to take the best checkpoint file from the last run, and train with more imagery/labels, in or to "fine tune" the model.
I tried:

!python run.py name=buildingseg eval_mode=True trainer.resume_from_checkpoint=/content/drive/MyDrive/ML5/LandCoverSeg/outputs/2024-10-21/22-33-24/checkpoints/epoch_120.ckpt ~logger

But it eventually just stops, even though its running with a Nvidia T4 on colab, with no errors:

[2024-10-30 23:15:11,210][src.train][INFO] - Logging hyperparameters!
[2024-10-30 23:15:11,213][src.train][INFO] - Starting training!
[2024-10-30 23:15:11,213][src.train][INFO] - Starting testing!
[2024-10-30 23:15:16,846][pytorch_lightning.utilities.rank_zero][INFO] - Restoring states from the checkpoint path at /content/drive/MyDrive/ML5/LandCoverSeg/outputs/2024-10-21/22-33-24/checkpoints/epoch_120.ckpt
[2024-10-30 23:15:34,289][pytorch_lightning.accelerators.gpu][INFO] - LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
[2024-10-30 23:15:34,290][pytorch_lightning.utilities.rank_zero][INFO] - Loaded model weights from checkpoint at /content/drive/MyDrive/ML5/LandCoverSeg/outputs/2024-10-21/22-33-24/checkpoints/epoch_120.ckpt
Testing DataLoader 0: 100% 17/17 [00:42<00:00,  2.50s/it]
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃        Test metric        ┃       DataLoader 0        ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│         test_dice         │    0.9988545775413513     │
│         test_iou          │    0.9977195858955383     │
│         test_loss         │    0.6067041158676147     │
└───────────────────────────┴───────────────────────────┘
[2024-10-30 23:16:17,053][src.train][INFO] - Finalizing!

How do I finetune an existing model with new imagery/labels?
Is there something in the command or configs I need to change?
Thank you for your thoughts.

[bartoszptak]

Hi @jducnuigeen,
The command has the eval_mode=True parameter, which indicates that you are loading a checkpoint and only evaluating it. Try removing the parameter and checking again.

[jducnuigeen]

Great, thank you. I'll try it without it, and that should be the last item to review before I submit my tutortial on my workflow.

[jducnuigeen]

An update after 2 months of struggle: While I got the first run above to work in Google Colab (which uses Tesla T4 GPU's using Cuda version 12.2 and the Turing architecture) I tried the same workflow on a local machine using a newer Nvidia RTX 2000ADA GPU which uses Cuda 12.4 and the ADA Lovelace architecture. Apparently these GPU's use different architectures and it makes a difference to the pytorch version. My trials on this newer GPU fails with lots of problems related to this GPU where it seems the torch kernel I used above has to be configured/compiled to use a this different GPU. I tried updating the torch version to 2.6.0 and the pytorch-lightning version 2.5.0, but now I get more error like

pytorch_lightning.utilities.exceptions.MisconfigurationException: The provided lr scheduler `CosineAnnealingLR` doesn't follow PyTorch's LRScheduler API. You should override the `LightningModule.lr_scheduler_step` hook with your own logic if you are using a custom LR scheduler.

So the tutorial I wanted to write will have to be on hold until I can figure out this rabbit hole of issues to solve.

Note: I stopped using Google Colab because it has a 12hr runtime limit.

[przemyslaw-aszkowski]

Hi,
Thanks for the update. Unfortunately, it is not easy to get everything running with machine learning, For the error you are seeing maybe it would help to use newer/older PyTorch and lightning versions. Or a different lr_scheduller.