update and add RI-GAN models and scripts

Author: MatthijsMars

configs/radio_meerkat_macro.yaml

@@ -0,0 +1,39 @@
+#Change checkpoint and sense_map path
+checkpoint_dir: /share/gpu0/mars/TNG_data/rcGAN/models/meerkat_macro/
+data_path: /share/gpu0/mars/TNG_data/rcGAN/meerkat_clean/
+
+# Define the experience
+experience: radio
+
+# Number of code vectors for each phase
+num_z_test: 32
+num_z_valid: 8
+num_z_train: 2
+
+# Data
+in_chans: 2  # Real+Imag parts from obs
+out_chans: 1
+im_size: 360 #384x384 pixel images
+
+# Options
+alt_upsample: False # False -> convt upsampling, True -> interpolate upsampling
+norm: macro # none, micro, macro
+
+# Optimizer:
+lr: 0.001
+beta_1: 0
+beta_2: 0.99
+
+# Loss weights
+gp_weight: 10
+adv_weight: 1e-5
+
+# Training
+batch_size: 2 # per GPU
+accumulate_grad_batches: 2
+
+#Remember to increase this for full training
+num_epochs: 100
+psnr_gain_tol: 0.25
+
+num_workers: 4

configs/radio_meerkat_macro_gradient.yaml

@@ -0,0 +1,40 @@
+#Change checkpoint and sense_map path
+checkpoint_dir: /share/gpu0/mars/TNG_data/rcGAN/models/meerkat_macro/
+data_path: /share/gpu0/mars/TNG_data/rcGAN/meerkat_clean/
+
+# Define the experience
+experience: radio
+
+# Number of code vectors for each phase
+num_z_test: 32
+num_z_valid: 8
+num_z_train: 2
+
+# Data
+in_chans: 2  # Real+Imag parts from obs
+out_chans: 1
+im_size: 360 #384x384 pixel images
+
+# Options
+alt_upsample: False # False -> convt upsampling, True -> interpolate upsampling
+norm: macro # none, micro, macro
+gradient: True
+
+# Optimizer:
+lr: 0.001
+beta_1: 0
+beta_2: 0.99
+
+# Loss weights
+gp_weight: 10
+adv_weight: 1e-5
+
+# Training
+batch_size: 2 # per GPU
+accumulate_grad_batches: 2
+
+#Remember to increase this for full training
+num_epochs: 100
+psnr_gain_tol: 0.25
+
+num_workers: 4

data/datasets/Radio_data.py

@@ -30,9 +30,17 @@ def __init__(self, data_dir, transform, norm='micro'):
             # if micro we do the normalisation in the transform
             pass
         elif norm == 'macro':
-            self.transform.mean_x, self.transform.std_x = self.x.mean(), self.x.std()
-            self.transform.mean_y, self.transform.std_y = self.y.mean(), self.y.std()
-            self.transform.mean_uv, self.transform.std_uv = self.uv.mean(), self.uv.std()
+            # load means and stds from train set
+            self.transform.mean_x  = np.load(data_dir.parent.joinpath("train/mean_x.npy"))
+            self.transform.std_x   = np.load(data_dir.parent.joinpath("train/std_x.npy"))
+            self.transform.mean_y  = np.load(data_dir.parent.joinpath("train/mean_y.npy"))
+            self.transform.std_y   = np.load(data_dir.parent.joinpath("train/std_y.npy"))
+            self.transform.mean_uv = np.load(data_dir.parent.joinpath("train/mean_uv.npy"))
+            self.transform.std_uv  = np.load(data_dir.parent.joinpath("train/std_uv.npy"))
+            
+#             self.transform.mean_x, self.transform.std_x = self.x.mean(), self.x.std()
+#             self.transform.mean_y, self.transform.std_y = self.y.mean(), self.y.std()
+#             self.transform.mean_uv, self.transform.std_uv = self.uv.mean(), self.uv.std()
 
     def __len__(self):
         """Returns the number of samples in the dataset."""
@@ -72,9 +80,17 @@ def __init__(self, data_dir, transform, norm='micro'):
             # if micro we do the normalisation in the transform
             pass
         elif norm == 'macro':
-            self.transform.mean_x, self.transform.std_x = self.x.mean(), self.x.std()
-            self.transform.mean_y, self.transform.std_y = self.y.mean(), self.y.std()
-            self.transform.mean_uv, self.transform.std_uv = self.uv.mean(), self.uv.std()
+            # load means and stds from train set
+            self.transform.mean_x  = np.load(data_dir.parent.joinpath("train/mean_x.npy"))
+            self.transform.std_x   = np.load(data_dir.parent.joinpath("train/std_x.npy"))
+            self.transform.mean_y  = np.load(data_dir.parent.joinpath("train/mean_y.npy"))
+            self.transform.std_y   = np.load(data_dir.parent.joinpath("train/std_y.npy"))
+            self.transform.mean_uv = np.load(data_dir.parent.joinpath("train/mean_uv.npy"))
+            self.transform.std_uv  = np.load(data_dir.parent.joinpath("train/std_uv.npy"))
+            
+#             self.transform.mean_x, self.transform.std_x = self.x.mean(), self.x.std()
+#             self.transform.mean_y, self.transform.std_y = self.y.mean(), self.y.std()
+#             self.transform.mean_uv, self.transform.std_uv = self.uv.mean(), self.uv.std()
 
     def __len__(self):
         """Returns the number of samples in the dataset."""
@@ -116,6 +132,16 @@ def __init__(self, data_dir, transform, norm='micro'):
             self.transform.mean_x, self.transform.std_x = self.x.mean(), np.mean(self.x.std(axis=(1,2)))
             self.transform.mean_y, self.transform.std_y = self.y.mean(), np.mean(self.y.std(axis=(1,2)))
             self.transform.mean_uv, self.transform.std_uv = self.uv.mean(),  np.mean(self.uv.std(axis=(1,2)))
+            
+            np.save(data_dir.joinpath("mean_x.npy"), self.transform.mean_x)
+            np.save(data_dir.joinpath("std_x.npy"), self.transform.std_x)
+            np.save(data_dir.joinpath("mean_y.npy"), self.transform.mean_y)
+            np.save(data_dir.joinpath("std_y.npy"), self.transform.std_y)
+            np.save(data_dir.joinpath("mean_uv.npy"), self.transform.mean_uv)
+            np.save(data_dir.joinpath("std_uv.npy"), self.transform.std_uv)
+            
+            
+            
 
     def __len__(self):
         """Returns the number of samples in the dataset."""