Re-architecting S2S2Net to be simpler (single-branch instead of dual-branch) and peform better (using a nicer loss model). Recommendations from @Ziwei-0129.

TODO:

• [x] Change from using dual-branch (in 979358749bdbca1587322ff4cd14504e882aa3b3) to single-branch (bf4f34efc46c9a01b4d9e79ea7571a1ec6abf15f)
  • Mainly because the dual-branch didn't really help for our task as it needed too much tweaking/hyperparameter tuning
  • Really, a single-branch network that's well designed can do the segmentation+super-resolution task equally well
• [x] Swap out Focal Loss for a BinaryCrossEntropyLoss+DiceLoss (from https://smp.readthedocs.io/en/v0.2.1/losses.html#diceloss), also using soft labels of 0.1 (28892798b3ad315e7ee1a6d77897cc3910f328f9)
• [x] Add Mish activation function to the upsampling head, change Segformer activation to Swish (f5da3c7a931432fe535de14a1837de80b0d5974f)
• [x] Change from x4x5 head to progressive upsampling head x2x2x2x2.5 (50dcc4a2a3870237854d8f9c25b953e1fb1170a4)

References:

• https://medium.com/ai-salon/understanding-dice-loss-for-crisp-boundary-detection-bb30c2e5f62b
• Deng, R., Shen, C., Liu, S., Wang, H., & Liu, X. (2018). Learning to predict crisp boundaries (arXiv:1807.10097). arXiv. https://doi.org/10.48550/arXiv.1807.10097
• Ma, J., Chen, J., Ng, M., Huang, R., Li, Y., Li, C., Yang, X., & Martel, A. L. (2021). Loss odyssey in medical image segmentation. Medical Image Analysis, 71, 102035. https://doi.org/10.1016/j.media.2021.102035
• Müller, R., Kornblith, S., & Hinton, G. (2020). When Does Label Smoothing Help? (arXiv:1906.02629). arXiv. https://doi.org/10.48550/arXiv.1906.02629

To prevent out-of-memory (OOM) errors when running the Transformer models! Change from distributed data parallel (DDP) to a data+model parallel strategy

Current State

As of 979358749bdbca1587322ff4cd14504e882aa3b3, we have been using distributed data parallel (DDP) to split the data batch-wise across multiple GPUs. However when running on a full-size Sentinel-2 image (batch_size=1) during test phase (#1), this can already cause out-of-memory issues for our Super-Resolution Segmentation task.

Future State

One possible solution is to shard the neural network model itself across multiple GPUs. This reduces the GPU memory requirements and allows for larger models and/or bigger datasets to be used for training/inference.

Specifically, we'll be switching to use DeepSpeed (https://github.com/microsoft/DeepSpeed) which offers several 'levels' of model sharding, and . See https://devblog.pytorchlightning.ai/experiment-with-billion-parameter-models-faster-using-deepspeed-and-meta-tensors-2e9c255edd71 and https://huggingface.co/blog/zero-deepspeed-fairscale for a good explainer

Main DeepSpeed stages (from https://pytorch-lightning.readthedocs.io/en/1.6.3/advanced/model_parallel.html#deepspeed):

• DeepSpeed ZeRO Stage 1 - Shard optimizer states, remains at speed parity with DDP whilst providing memory improvement
• DeepSpeed ZeRO Stage 2 - Shard optimizer states and gradients, remains at speed parity with DDP whilst providing even more memory improvement
• DeepSpeed ZeRO Stage 3 - Shard optimizer states, gradients, parameters and optionally activations. Increases distributed communication volume, but provides even more memory improvement

:bulb: Suggest to use Stage 2 instead of Stage 3 because while Stage 3 improves memory use, it comes with increased latency from the cost of extra distributed communication.

Other benefits of using DeepSpeed:

• Stage 2 and Stage 3 also has an 'Offload' to CPU feature to save on memory, in cases when the GPU memory is simply not enough
• Allows me to train the model on just 16GB of GPU RAM on my workstation :exploding_head:

Alternative strategies (and why they were not considered)

Pytorch-Lightning offers several other advanced training strategies. These might work well for other cases, but probably not for our specific project.

• Bagua (https://github.com/BaguaSys/bagua)
  • Data Distributed Parallel
  • Why not use this? Due to it not being model parallel (though they may be working on it)
  • https://pytorch-lightning.readthedocs.io/en/1.6.3/accelerators/gpu.html#bagua
  • https://devblog.pytorchlightning.ai/bagua-a-new-efficient-distributed-training-strategy-available-in-pytorch-lightning-1-6-d6392633b15
• Fairscale (https://github.com/facebookresearch/fairscale)
  • Model parallel training, a close competitor to DeepSpeed
  • Why not use this? I did try, but the conda-forge package couldn't work because of some ABI compatibility issue.
  • https://pytorch-lightning.readthedocs.io/en/1.6.3/advanced/model_parallel.html#fully-sharded-training

TODO:

• [x] Add deepspeed dependency (0a666012ffd523f8fffed541612d85ed0e151e7d)
• [x] Switch model to use DeepSpeed ZeRO Stage 2 (6394c12fc92a4837f8f453b5192d038928dffd46)
• [ ] ~~Use Meta Tensors, c.f. https://devblog.pytorchlightning.ai/experiment-with-billion-parameter-models-faster-using-deepspeed-and-meta-tensors-2e9c255edd71~~
  • Nope, doesn't work. Error given is NotImplementedError: Could not run 'aten::_local_scalar_dense' with arguments from the 'Meta' backend. See also https://github.com/pytorch/pytorch/issues/77764
• [ ] Decide whether to remove the Super-Resolution branch :thinking:

To allow fair comparison of different models, compute F1 and IoU metrics on 8 independent images from a hold-out set (see 00a1c5b4fa9992749221d3f02bbd8a1e99447890) that was handpicked to be from different geographic localities compared to the train/val set. The evaluate function now has a boolean calc_loss parameter, so that the expensive loss computation can be skipped during the test_step loop (i.e. compute metrics only).

Note on issues and workarounds:

• Computing the test metrics means the full-size masks need to be loaded, and for some reason the geographical extent of the mask doesn't align with the image so need to do rio.clip_box.
• Predicted tensor shapes are sometimes not exactly 5x that of the input image, so need to interpolate to the right size (which isn't actually proper, I know).
• ~~Needed to use CPU for test_step because GPU did not have enough memory, and use 32-bit precision too.~~ Edit: fixed with model sharding #2.

TODO:

• [x] Initial implementatin of test_step (7e87c6ccb0ded0586b90a338d0be374b0af1737a)
• [x] Run test_step on GPU instead of CPU (#1 and 099f92f5bb441a0158b51230a1108f38cd4cdbd0)
• [x] Refactor to ensure image and mask geographical extents align (e03a6e97f85b4ceca2f25fba84d77e3087820d15)
• [ ] FIx saving and loading from model checkpoint