SimVPv2 모델 수정 사항 체크 요청

.gitignore

@@ -1,5 +1,7 @@
 __pycache__
 /results
+exp_results/
 /figs
 .vscode
-.ipynb_checkpoints
+.ipynb_checkpoints
+.DS_Store

API/dataloader_s4a.py

@@ -35,11 +35,12 @@
 
 NORMALIZATION_DIV = 10000
 
+
 def min_max_normalize(image, percentile=2):
     image = image.astype('float32')
 
-    percent_min = np.percentile(image, percentile, axis=(0,1))
-    percent_max = np.percentile(image, 100-percentile, axis=(0,1))
+    percent_min = np.percentile(image, percentile, axis=(0, 1))
+    percent_max = np.percentile(image, 100-percentile, axis=(0, 1))
 
     mask = np.mean(image, axis=2) != 0
     if image.shape[1] * image.shape[0] - np.sum(mask) > 0:
@@ -48,9 +49,9 @@ def min_max_normalize(image, percentile=2):
         percent_min = np.nanpercentile(mdata, percentile, axis=(0, 1))
 
     norm = (image-percent_min) / (percent_max - percent_min)
-    norm[norm<0] = 0
-    norm[norm>1] = 1
-    norm = norm * mask[:,:,np.newaxis]
+    norm[norm < 0] = 0
+    norm[norm > 1] = 1
+    norm = norm * mask[:, :, np.newaxis]
     # norm = (norm * 255).astype('uint8') * mask[:,:,np.newaxis]
 
     return norm
@@ -117,12 +118,12 @@ def __init__(
             root_dir: str = None,
             img_dir: str = None,
             ann_dir: str = None,
-            band_mode: str ='nrgb',
+            band_mode: str = 'nrgb',
             bands: list = None,
             linear_encoder: dict = None,
             start_month: int = 0,
             end_month: int = 12,
-            output_size: tuple = (64,64),
+            output_size: tuple = (64, 64),
             binary_labels: bool = False,
             return_parcels: bool = False,
             mode: str = 'test',
@@ -154,8 +155,9 @@ def __init__(
             self.bands = ['B02', 'B03', 'B04', 'B08']
 
         elif band_mode == 'rdeg':
-            self.bands = ['B02', 'B03', 'B04', 'B08', 'B05', 'B06', 'B07', 'B8A', 'B11', 'B12']
-
+            self.bands = ['B02', 'B03', 'B04', 'B08',
+                          'B05', 'B06', 'B07', 'B8A', 'B11', 'B12']
+
         else:
             raise RuntimeError
 
@@ -204,10 +206,12 @@ def __init__(
         self.std = 1
 
         print('Rootdir: {}'.format(self.root_dir))
-        print('Scenario: {}, MODE: {}, length of datasets: {}'.format(self.scenario, self.mode, len(self.img_infos)))
-        print('Acquired Data Month: From {} to {}'.format(self.start_month + 1, self.end_month + 1))
-        print(f'Data shape: [T, C, H, W] ({self.end_month - self.start_month}, {len(self.bands)}, {output_size[0]}, {output_size[0]})')
-
+        print('Scenario: {}, MODE: {}, length of datasets: {}'.format(
+            self.scenario, self.mode, len(self.img_infos)))
+        print('Acquired Data Month: From {} to {}'.format(
+            self.start_month + 1, self.end_month + 1))
+        print(
+            f'Data shape: [T, C, H, W] ({self.end_month - self.start_month}, {len(self.bands)}, {output_size[0]}, {output_size[0]})')
 
     def prepare_train_img(self, idx: int) -> dict:
         if self.band_mode == 'nrgb':
@@ -217,7 +221,8 @@ def prepare_train_img(self, idx: int) -> dict:
             readpath = os.path.join(self.img_dir, self.img_infos[idx])
             img = np.load(readpath)
 
-            rdegpath = os.path.join(os.path.dirname(self.img_dir), 'rdeg', self.img_infos[idx])
+            rdegpath = os.path.join(os.path.dirname(
+                self.img_dir), 'rdeg', self.img_infos[idx])
             rdeg = np.load(rdegpath)
             img = np.stack([img, rdeg], axis=1)
         else:
@@ -237,35 +242,35 @@ def _normalize(self, img):
         if self.min_max_normalize:
             T, C, H, W = img.shape
             img = img.reshape(T*C, H, W)
-            img = min_max_normalize(img.transpose(1,2,0), percentile=0.5)
-            img = img.transpose(2,0,1)
+            img = min_max_normalize(img.transpose(1, 2, 0), percentile=0.5)
+            img = img.transpose(2, 0, 1)
             img = img.reshape(T, C, H, W)
         else:
-            img = np.divide(img, NORMALIZATION_DIV) #  / 10000
+            img = np.divide(img, NORMALIZATION_DIV)  # / 10000
         return img
 
-
     def __getitem__(self, idx: int) -> dict:
         img, ann = self.prepare_train_img(idx)
 
         # Normalize data to range [0-1]
         img = self._normalize(img)
 
-        # out = {}
-        # if self.return_parcels:
-        #     parcels = ann != 0
-        #     out['parcels'] = parcels
-
-        # if self.binary_labels:
-        #     # Map 0: background class, 1: parcel
-        #     ann[ann != 0] = 1
-        # else:
-        #     # Map labels to 0-len(unique(crop_id)) see config
-        #     # labels = np.vectorize(self.linear_encoder.get)(labels)
-        #     _ = np.zeros_like(ann)
-        #     for crop_id, linear_id in self.linear_encoder.items():
-        #         _[ann == crop_id] = linear_id
-        #     ann = _
+    ########TODO########
+        out = {}
+        if self.return_parcels:
+            parcels = ann != 0
+            out['parcels'] = parcels
+
+        if self.binary_labels:
+            # Map 0: background class, 1: parcel
+            ann[ann != 0] = 1
+        else:
+            # Map labels to 0-len(unique(crop_id)) see config
+            # labels = np.vectorize(self.linear_encoder.get)(labels)
+            _ = np.zeros_like(ann)
+            for crop_id, linear_id in self.linear_encoder.items():
+                _[ann == crop_id] = linear_id
+            ann = _
 
         # # # Map all classes NOT in linear encoder's values to 0
         # ann[~np.isin(ann, list(self.linear_encoder.values()))] = 0
@@ -277,7 +282,11 @@ def __getitem__(self, idx: int) -> dict:
         input = torch.from_numpy(input).contiguous().float()
 
         if self.get_ann:
-            ann = ann.astype('int16')
+            ann = ann.astype('float32')
+            _ = np.zeros_like(ann)
+            for crop_id, linear_id in LINEAR_ENCODER.items():
+                _[ann == crop_id] = linear_id
+            ann = _ / len(SELECTED_CLASSES)
             ann = torch.from_numpy(ann).contiguous().long()
             return input, output, ann
         else:
@@ -288,10 +297,11 @@ def __len__(self):
 
 
 def load_data(batch_size, val_batch_size, num_workers, data_root):
-
-    train_set = NpyPADDataset(root_dir=data_root, band_mode='nrgb', start_month=1, end_month=13, mode='train')
-    test_set = NpyPADDataset(root_dir=data_root, band_mode='nrgb', start_month=1, end_month=13, mode='val')
-
+    # get_ann 옵션 추가
+    train_set = NpyPADDataset(
+        root_dir=data_root, band_mode='nrgb', start_month=1, end_month=13, get_ann=False, mode='train')
+    test_set = NpyPADDataset(
+        root_dir=data_root, band_mode='nrgb', start_month=1, end_month=13, get_ann=False, mode='val')
 
     dataloader_train = torch.utils.data.DataLoader(
         train_set, batch_size=batch_size, shuffle=True, pin_memory=True, drop_last=True, num_workers=num_workers)
@@ -305,5 +315,10 @@ def load_data(batch_size, val_batch_size, num_workers, data_root):
 
 if __name__ == '__main__':
     rootdir = '/home/jovyan/shared_volume/data/newdata'
-    dataset =  NpyPADDataset(root_dir=rootdir, band_mode='nrgb', start_month=1, end_month=13, mode='val')
-    print(dataset[0])
+    dataset = NpyPADDataset(
+        root_dir=rootdir, band_mode='nrgb', start_month=1, end_month=13, get_ann=True, mode='val')
+
+    dataloader = torch.utils.data.DataLoader(
+        dataset, batch_size=batch_size, shuffle=True, pin_memory=True, drop_last=True, num_workers=num_workers)
+
+    # print(dataset[0][2])

configs/SimVP.py

@@ -1,6 +1,6 @@
 method = 'SimVP'
 model_type = 'gSTA'
-hid_S = 64
-hid_T = 512
+hid_S = 16
+hid_T = 256
 N_T = 8
-N_S = 4
+N_S = 4

constants/dataset_constant.py

@@ -6,6 +6,9 @@
         'total_length': 20
     },
     's4a': {
+        # in_shape 수정
+        # [8, 4, 64, 64] -> [9, 4, 64, 64]
+        # 'in_shape': [9, 4, 64, 64],
         'in_shape': [8, 4, 64, 64],
         'pre_seq_length': 8,
         'aft_seq_length': 4,
@@ -47,4 +50,4 @@
         'aft_seq_length': 12,
         'total_length': 24
     }
-}
+}

main.py

@@ -51,6 +51,7 @@ def __init__(self, args):
             _tmp_input = torch.ones(1, self.args.total_length, Hp, Wp, Cp).to(self.device)
             _tmp_flag = torch.ones(1, self.args.total_length - 2, Hp, Wp, Cp).to(self.device)
             flops = FlopCountAnalysis(self.method.model, (_tmp_input, _tmp_flag))
+        '''RuntimeError 발생'''
         print_log(flop_count_table(flops))
 
     def _acquire_device(self):
@@ -112,6 +113,20 @@ def train(self):
         num_updates = 0
         # constants for other methods:
         eta = 1.0  # PredRNN
+        ##### for size test #####
+        # train_inputs, train_outputs, train_labels,  = next(iter(self.train_loader))
+        # print(f'Inputs batch shape: {train_inputs.size()}')
+        # print(f'Outputs batch shape: {train_outputs.size()}')
+        # print(f'Labels batch shape: {train_labels.size()}')
+        # _label = np.array(train_labels)
+        # label = torch.from_numpy(_label)
+        # label = label.unsqueeze(1)
+        # conv = torch.nn.Conv2d(in_channels=1, out_channels=4, kernel_size=1)
+        # label = conv(label.float())
+        # print(f'NEW Labels batch shape: {label.size()}')
+
+        # quit()
+        #########################
         for epoch in range(self.config['epoch']):
             loss_mean = 0.0
 

methods/simvp.py

@@ -14,6 +14,7 @@ def __init__(self, args, device, steps_per_epoch):
         self.model = self._build_model(self.config)
         self.model_optim, self.scheduler = self._init_optimizer(steps_per_epoch)
         self.criterion = nn.MSELoss()
+        # self.conv_c4 = nn.Conv2d(in_channels=1, out_channels=4, kernel_size=1)
 
     def _build_model(self, config):
         return SimVP_Model(**config).to(self.device)
@@ -46,8 +47,27 @@ def train_one_epoch(self, train_loader, epoch, num_updates, loss_mean, **kwargs)
         self.model.train()
 
         train_pbar = tqdm(train_loader)
+        # ann 추출 후 batch_x에 추가
+        # for batch_x, batch_y, ann in train_pbar:
         for batch_x, batch_y in train_pbar:
             self.model_optim.zero_grad()
+            # [16, 64, 64] -> [16, 1, 64, 64]
+            # B H W -> B C H W
+            # ann = ann.unsqueeze(1)
+
+            # [16, 1, 64, 64] -> [16, 4, 64, 64]
+            # ann = self.conv_c4(ann.float())
+
+            # repeat으로 차원 추가
+            # ann = ann.repeat(1,4,1,1)
+
+            # [16, 4, 64, 64] -> [16, 1, 4, 64, 64]
+            # B C H W -> B T C H W
+            # ann = ann.unsqueeze(1)
+
+            # [16, 8, 4, 64, 64] -> [16, 9, 4, 64, 64]
+            # batch_x = torch.cat([batch_x, ann], dim=1)
+
             batch_x, batch_y = batch_x.to(self.device), batch_y.to(self.device)
             pred_y = self._predict(batch_x)
 
@@ -68,7 +88,14 @@ def vali_one_epoch(self, vali_loader, **kwargs):
         self.model.eval()
         preds_lst, trues_lst, total_loss = [], [], []
         vali_pbar = tqdm(vali_loader)
+        # for i, (batch_x, batch_y, ann) in enumerate(vali_pbar):
         for i, (batch_x, batch_y) in enumerate(vali_pbar):
+            # ann = ann.unsqueeze(1)
+            # ann = self.conv_c4(ann.float())
+            # ann = ann.repeat(1,4,1,1)
+
+            # ann = ann.unsqueeze(1)
+            # batch_x = torch.cat([batch_x, ann], dim=1)
             batch_x, batch_y = batch_x.to(self.device), batch_y.to(self.device)
             pred_y = self._predict(batch_x)
             loss = self.criterion(pred_y, batch_y)
@@ -91,7 +118,14 @@ def test_one_epoch(self, test_loader, **kwargs):
         self.model.eval()
         inputs_lst, trues_lst, preds_lst = [], [], []
         test_pbar = tqdm(test_loader)
+        # for batch_x, batch_y, ann in test_pbar:
         for batch_x, batch_y in test_pbar:
+            # ann = ann.unsqueeze(1)
+            # ann = self.conv_c4(ann.float())
+            # ann = ann.repeat(1,4,1,1)
+
+            # ann = ann.unsqueeze(1)
+            # batch_x = torch.cat([batch_x, ann], dim=1)
             pred_y = self._predict(batch_x.to(self.device))
 
             list(map(lambda data, lst: lst.append(data.detach().cpu().numpy()), [