Update burn_scars_Prithvi_100M.py

burn_scars_Prithvi_100M.py

@@ -1,11 +1,17 @@
+import os 
+custom_imports = dict(imports=['geospatial_fm'])
+
 dist_params = dict(backend='nccl')
 log_level = 'INFO'
 load_from = None
 resume_from = None
 cudnn_benchmark = True
-custom_imports = dict(imports=['geospatial_fm'])
+
 dataset_type = 'GeospatialDataset'
-data_root = '/dccstor/geofm-finetuning/fire-scars/finetune-data/6_bands_no_replant_extended'
+
+# TO BE DEFINED BY USER: data directory
+data_root = '<path to data root>'
+
 num_frames = 1
 img_size = 224
 num_workers = 4
@@ -22,45 +28,45 @@ img_norm_cfg = dict(
 bands = [0, 1, 2, 3, 4, 5]
 tile_size = 224
 orig_nsize = 512
-crop_size = (224, 224)
+crop_size = (tile_size, tile_size)
 img_suffix = '_merged.tif'
 seg_map_suffix = '.mask.tif'
 ignore_index = -1
 image_nodata = -9999
 image_nodata_replace = 0
 image_to_float32 = True
-# pretrained_weights_path = '/dccstor/geofm-finetuning/pretrain_ckpts/mae_weights/2023-04-29_21-50-47/epoch-725-loss-0.0365.pt'
-pretrained_weights_path = None
+
+# model
+# TO BE DEFINED BY USER: model path
+pretrained_weights_path = '<path to pretrained weights>'
 num_layers = 12
 patch_size = 16
 embed_dim = 768
 num_heads = 12
 tubelet_size = 1
-epochs = 50
-eval_epoch_interval = 5
-experiment = 'test2'
-project_dir = '/dccstor/geofm-finetuning/fire-scars/os'
-work_dir = '/dccstor/geofm-finetuning/fire-scars/os/test2'
-save_path = '/dccstor/geofm-finetuning/fire-scars/os/test2'
+output_embed_dim = num_frames*embed_dim
+max_intervals=10000
+evaluation_interval=1000
+
+# TO BE DEFINED BY USER: model path
+experiment = '<experiment name>'
+project_dir = '<project directory name>'
+work_dir = os.path.join(project_dir, experiment)
+save_path = work_dir
+
+save_path = work_dir
 train_pipeline = [
-    dict(type='LoadGeospatialImageFromFile', to_float32=True),
+    dict(type='LoadGeospatialImageFromFile', to_float32=image_to_float32, channels_last=True),
     dict(type='LoadGeospatialAnnotations', reduce_zero_label=False),
-    dict(type='BandsExtract', bands=[0, 1, 2, 3, 4, 5]),
+    dict(type='BandsExtract', bands=bands),
     dict(type='RandomFlip', prob=0.5),
     dict(type='ToTensor', keys=['img', 'gt_semantic_seg']),
-    dict(
-        type='TorchNormalize',
-        means=[
-            0.033349706741586264, 0.05701185520536176, 0.05889748132001316,
-            0.2323245113436119, 0.1972854853760658, 0.11944914225186566
-        ],
-        stds=[
-            0.02269135568823774, 0.026807560223070237, 0.04004109844362779,
-            0.07791732423672691, 0.08708738838140137, 0.07241979477437814
-        ]),
-    dict(type='TorchRandomCrop', crop_size=(224, 224)),
-    dict(type='Reshape', keys=['img'], new_shape=(6, 1, 224, 224)),
-    dict(type='Reshape', keys=['gt_semantic_seg'], new_shape=(1, 224, 224)),
+    # to channels first
+    dict(type="TorchPermute", keys=["img"], order=(2, 0, 1)),
+    dict(type='TorchNormalize', **img_norm_cfg),
+    dict(type='TorchRandomCrop', crop_size=(tile_size, tile_size)),
+    dict(type='Reshape', keys=['img'], new_shape=(len(bands), num_frames, tile_size, tile_size)),
+    dict(type='Reshape', keys=['gt_semantic_seg'], new_shape=(1, tile_size, tile_size)),
     dict(
         type='CastTensor',
         keys=['gt_semantic_seg'],
@@ -68,23 +74,16 @@ train_pipeline = [
     dict(type='Collect', keys=['img', 'gt_semantic_seg'])
 ]
 test_pipeline = [
-    dict(type='LoadGeospatialImageFromFile', to_float32=True),
-    dict(type='BandsExtract', bands=[0, 1, 2, 3, 4, 5]),
+    dict(type='LoadGeospatialImageFromFile', to_float32=image_to_float32, channels_last=True),
+    dict(type='BandsExtract', bands=bands),
     dict(type='ToTensor', keys=['img']),
-    dict(
-        type='TorchNormalize',
-        means=[
-            0.033349706741586264, 0.05701185520536176, 0.05889748132001316,
-            0.2323245113436119, 0.1972854853760658, 0.11944914225186566
-        ],
-        stds=[
-            0.02269135568823774, 0.026807560223070237, 0.04004109844362779,
-            0.07791732423672691, 0.08708738838140137, 0.07241979477437814
-        ]),
+    # to channels first
+    dict(type="TorchPermute", keys=["img"], order=(2, 0, 1)),
+    dict(type='TorchNormalize', **img_norm_cfg),
     dict(
         type='Reshape',
         keys=['img'],
-        new_shape=(6, 1, -1, -1),
+        new_shape=(len(bands), num_frames, -1, -1),
         look_up=dict({
             '2': 1,
             '3': 2
@@ -99,136 +98,43 @@ test_pipeline = [
             'scale_factor', 'img_norm_cfg'
         ])
 ]
+
+CLASSES = ('Unburnt land', 'Burn scar')
+
 data = dict(
-    samples_per_gpu=4,
-    workers_per_gpu=4,
+    samples_per_gpu=samples_per_gpu,
+    workers_per_gpu=num_workers,
     train=dict(
-        type='FireScars',
-        data_root=
-        '/dccstor/geofm-finetuning/fire-scars/finetune-data/6_bands_no_replant_extended',
+        type=dataset_type,
+        CLASSES=CLASSES,
+        data_root=data_root,
         img_dir='training',
         ann_dir='training',
-        img_suffix='_merged.tif',
-        seg_map_suffix='.mask.tif',
-        pipeline=[
-            dict(type='LoadGeospatialImageFromFile', to_float32=True),
-            dict(type='LoadGeospatialAnnotations', reduce_zero_label=False),
-            dict(type='BandsExtract', bands=[0, 1, 2, 3, 4, 5]),
-            dict(type='RandomFlip', prob=0.5),
-            dict(type='ToTensor', keys=['img', 'gt_semantic_seg']),
-            dict(
-                type='TorchNormalize',
-                means=[
-                    0.033349706741586264, 0.05701185520536176,
-                    0.05889748132001316, 0.2323245113436119,
-                    0.1972854853760658, 0.11944914225186566
-                ],
-                stds=[
-                    0.02269135568823774, 0.026807560223070237,
-                    0.04004109844362779, 0.07791732423672691,
-                    0.08708738838140137, 0.07241979477437814
-                ]),
-            dict(type='TorchRandomCrop', crop_size=(224, 224)),
-            dict(type='Reshape', keys=['img'], new_shape=(6, 1, 224, 224)),
-            dict(
-                type='Reshape',
-                keys=['gt_semantic_seg'],
-                new_shape=(1, 224, 224)),
-            dict(
-                type='CastTensor',
-                keys=['gt_semantic_seg'],
-                new_type='torch.LongTensor'),
-            dict(type='Collect', keys=['img', 'gt_semantic_seg'])
-        ],
+        img_suffix=img_suffix,
+        seg_map_suffix=seg_map_suffix,
+        pipeline=train_pipeline,
         ignore_index=-1),
     val=dict(
-        type='FireScars',
-        data_root=
-        '/dccstor/geofm-finetuning/fire-scars/finetune-data/6_bands_no_replant_extended',
+        type=dataset_type,
+        CLASSES=CLASSES,
+        data_root=data_root,
         img_dir='validation',
         ann_dir='validation',
-        img_suffix='_merged.tif',
-        seg_map_suffix='.mask.tif',
-        pipeline=[
-            dict(type='LoadGeospatialImageFromFile', to_float32=True),
-            dict(type='BandsExtract', bands=[0, 1, 2, 3, 4, 5]),
-            dict(type='ToTensor', keys=['img']),
-            dict(
-                type='TorchNormalize',
-                means=[
-                    0.033349706741586264, 0.05701185520536176,
-                    0.05889748132001316, 0.2323245113436119,
-                    0.1972854853760658, 0.11944914225186566
-                ],
-                stds=[
-                    0.02269135568823774, 0.026807560223070237,
-                    0.04004109844362779, 0.07791732423672691,
-                    0.08708738838140137, 0.07241979477437814
-                ]),
-            dict(
-                type='Reshape',
-                keys=['img'],
-                new_shape=(6, 1, -1, -1),
-                look_up=dict({
-                    '2': 1,
-                    '3': 2
-                })),
-            dict(
-                type='CastTensor', keys=['img'], new_type='torch.FloatTensor'),
-            dict(
-                type='CollectTestList',
-                keys=['img'],
-                meta_keys=[
-                    'img_info', 'seg_fields', 'img_prefix', 'seg_prefix',
-                    'filename', 'ori_filename', 'img', 'img_shape',
-                    'ori_shape', 'pad_shape', 'scale_factor', 'img_norm_cfg'
-                ])
-        ],
+        img_suffix=img_suffix,
+        seg_map_suffix=seg_map_suffix,
+        pipeline=test_pipeline,
         ignore_index=-1),
     test=dict(
-        type='FireScars',
-        data_root=
-        '/dccstor/geofm-finetuning/fire-scars/finetune-data/6_bands_no_replant_extended',
+        type=dataset_type,
+        CLASSES=CLASSES,
+        data_root=data_root,
         img_dir='validation',
         ann_dir='validation',
-        img_suffix='_merged.tif',
-        seg_map_suffix='.mask.tif',
-        pipeline=[
-            dict(type='LoadGeospatialImageFromFile', to_float32=True),
-            dict(type='BandsExtract', bands=[0, 1, 2, 3, 4, 5]),
-            dict(type='ToTensor', keys=['img']),
-            dict(
-                type='TorchNormalize',
-                means=[
-                    0.033349706741586264, 0.05701185520536176,
-                    0.05889748132001316, 0.2323245113436119,
-                    0.1972854853760658, 0.11944914225186566
-                ],
-                stds=[
-                    0.02269135568823774, 0.026807560223070237,
-                    0.04004109844362779, 0.07791732423672691,
-                    0.08708738838140137, 0.07241979477437814
-                ]),
-            dict(
-                type='Reshape',
-                keys=['img'],
-                new_shape=(6, 1, -1, -1),
-                look_up=dict({
-                    '2': 1,
-                    '3': 2
-                })),
-            dict(
-                type='CastTensor', keys=['img'], new_type='torch.FloatTensor'),
-            dict(
-                type='CollectTestList',
-                keys=['img'],
-                meta_keys=[
-                    'img_info', 'seg_fields', 'img_prefix', 'seg_prefix',
-                    'filename', 'ori_filename', 'img', 'img_shape',
-                    'ori_shape', 'pad_shape', 'scale_factor', 'img_norm_cfg'
-                ])
-        ],
+        img_suffix=img_suffix,
+        seg_map_suffix=seg_map_suffix,
+        pipeline=test_pipeline,
         ignore_index=-1))
+
 optimizer = dict(type='Adam', lr=1.3e-05, betas=(0.9, 0.999))
 optimizer_config = dict(grad_clip=None)
 lr_config = dict(
@@ -248,16 +154,20 @@ log_config = dict(
 checkpoint_config = dict(
     by_epoch=True,
     interval=10,
-    out_dir=
-    '/dccstor/geofm-finetuning/carlosgomes/fire_scars/carlos_replicate_experiment_fixed_lr'
+    out_dir=save_path 
 )
 evaluation = dict(
-    interval=1180,
+    interval=evaluation_interval,
     metric='mIoU',
     pre_eval=True,
     save_best='mIoU',
     by_epoch=False)
-runner = dict(type='IterBasedRunner', max_iters=6300)
+
+loss_func=dict(
+            type='DiceLoss', use_sigmoid=False, loss_weight=1,
+            ignore_index=-1)
+
+runner = dict(type='IterBasedRunner', max_iters=max_intervals)
 workflow = [('train', 1)]
 norm_cfg = dict(type='BN', requires_grad=True)
 model = dict(
@@ -265,28 +175,27 @@ model = dict(
     frozen_backbone=False,
     backbone=dict(
         type='TemporalViTEncoder',
-        pretrained=None,
-        # '/dccstor/geofm-finetuning/pretrain_ckpts/mae_weights/2023-04-29_21-50-47/epoch-725-loss-0.0365.pt',
-        img_size=224,
-        patch_size=16,
-        num_frames=1,
-        tubelet_size=1,
-        in_chans=6,
-        embed_dim=768,
+        pretrained=pretrained_weights_path,
+        img_size=img_size,
+        patch_size=patch_size,
+        num_frames=num_frames,
+        tubelet_size=tubelet_size,
+        in_chans=len(bands),
+        embed_dim=embed_dim,
         depth=12,
-        num_heads=12,
+        num_heads=num_heads,
         mlp_ratio=4.0,
         norm_pix_loss=False),
     neck=dict(
         type='ConvTransformerTokensToEmbeddingNeck',
-        embed_dim=768,
-        output_embed_dim=768,
+        embed_dim=embed_dim*num_frames,
+        output_embed_dim=output_embed_dim,
         drop_cls_token=True,
         Hp=14,
         Wp=14),
     decode_head=dict(
-        num_classes=2,
-        in_channels=768,
+        num_classes=len(CLASSES),
+        in_channels=output_embed_dim,
         type='FCNHead',
         in_index=-1,
         channels=256,
@@ -295,12 +204,11 @@ model = dict(
         dropout_ratio=0.1,
         norm_cfg=dict(type='BN', requires_grad=True),
         align_corners=False,
-        loss_decode=dict(
-            type='DiceLoss', use_sigmoid=False, loss_weight=1,
-            ignore_index=-1)),
+        loss_decode=
+        loss_decode=loss_func),
     auxiliary_head=dict(
-        num_classes=2,
-        in_channels=768,
+        num_classes=len(CLASSES),
+        in_channels=output_embed_dim,
         type='FCNHead',
         in_index=-1,
         channels=256,
@@ -309,10 +217,8 @@ model = dict(
         dropout_ratio=0.1,
         norm_cfg=dict(type='BN', requires_grad=True),
         align_corners=False,
-        loss_decode=dict(
-            type='DiceLoss', use_sigmoid=False, loss_weight=1,
-            ignore_index=-1)),
+        loss_decode=loss_func),
     train_cfg=dict(),
-    test_cfg=dict(mode='slide', stride=(112, 112), crop_size=(224, 224)))
+    test_cfg=dict(mode='slide', stride=(tile_size/2, tile_size/2), crop_size=(tile_size, tile_size)))
 gpu_ids = range(0, 1)
-auto_resume = False
+auto_resume = False