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.