Upload init

Author: hanyoseob

dataset.py

@@ -73,8 +73,6 @@ def __getitem__(self, index):
         if data.shape[0] > data.shape[1]:
             data = data.transpose((1, 0, 2))
 
-        sz = data.shape
-
         label = data + self.noise[index]
         input, mask = self.generate_mask(copy.deepcopy(label))
 

main.py

@@ -9,29 +9,40 @@
 cudnn.benchmark = True
 cudnn.fastest = True
 
+# FLAG_PLATFORM = 'laptop'
+FLAG_PLATFORM = 'colab'
+
 ## setup parse
 parser = argparse.ArgumentParser(description='Train the unet network',
                                  formatter_class=argparse.ArgumentDefaultsHelpFormatter)
 
-parser.add_argument('--gpu_ids', default='-1', dest='gpu_ids')
+if FLAG_PLATFORM == 'colab':
+    parser.add_argument('--gpu_ids', default='0', dest='gpu_ids')
 
-parser.add_argument('--mode', default='train', choices=['train', 'test'], dest='mode')
-parser.add_argument('--train_continue', default='off', choices=['on', 'off'], dest='train_continue')
+    parser.add_argument('--dir_checkpoint', default='./drive/My Drive/GitHub/pytorch-noise2void/checkpoints', dest='dir_checkpoint')
+    parser.add_argument('--dir_log', default='./drive/My Drive/GitHub/pytorch-noise2void/log', dest='dir_log')
+    parser.add_argument('--dir_result', default='./drive/My Drive/GitHub/pytorch-noise2void/results', dest='dir_result')
+    parser.add_argument('--dir_data', default='./drive/My Drive/datasets', dest='dir_data')
+elif FLAG_PLATFORM == 'laptop':
+    parser.add_argument('--gpu_ids', default='-1', dest='gpu_ids')
 
-parser.add_argument('--scope', default='resnet', dest='scope')
-parser.add_argument('--norm', type=str, default='inorm', dest='norm')
+    parser.add_argument('--dir_checkpoint', default='./checkpoints', dest='dir_checkpoint')
+    parser.add_argument('--dir_log', default='./log', dest='dir_log')
+    parser.add_argument('--dir_result', default='./results', dest='dir_result')
+    parser.add_argument('--dir_data', default='./datasets', dest='dir_data')
+
+parser.add_argument('--mode', default='train', choices=['train', 'test'], dest='mode')
+parser.add_argument('--train_continue', default='on', choices=['on', 'off'], dest='train_continue')
 
-parser.add_argument('--dir_checkpoint', default='./checkpoints', dest='dir_checkpoint')
-parser.add_argument('--dir_log', default='./log', dest='dir_log')
+parser.add_argument('--scope', default='denoising_resnet', dest='scope')
+parser.add_argument('--norm', type=str, default='bnorm', dest='norm')
 
 parser.add_argument('--name_data', type=str, default='bsd500', dest='name_data')
-parser.add_argument('--dir_data', default='../datasets', dest='dir_data')
-parser.add_argument('--dir_result', default='./results', dest='dir_result')
 
 parser.add_argument('--num_epoch', type=int,  default=300, dest='num_epoch')
-parser.add_argument('--batch_size', type=int, default=4, dest='batch_size')
+parser.add_argument('--batch_size', type=int, default=1, dest='batch_size')
 
-parser.add_argument('--lr_G', type=float, default=1e-4, dest='lr_G')
+parser.add_argument('--lr_G', type=float, default=1e-3, dest='lr_G')
 
 parser.add_argument('--optim', default='adam', choices=['sgd', 'adam', 'rmsprop'], dest='optim')
 parser.add_argument('--beta1', default=0.5, dest='beta1')
@@ -52,8 +63,8 @@
 
 parser.add_argument('--data_type', default='float32', dest='data_type')
 
-parser.add_argument('--num_freq_disp', type=int,  default=1, dest='num_freq_disp')
-parser.add_argument('--num_freq_save', type=int,  default=1, dest='num_freq_save')
+parser.add_argument('--num_freq_disp', type=int,  default=10, dest='num_freq_disp')
+parser.add_argument('--num_freq_save', type=int,  default=50, dest='num_freq_save')
 
 PARSER = Parser(parser)
 

train.py

@@ -146,8 +146,8 @@ def train(self):
         dataset_train = Dataset(dir_data_train, data_type=self.data_type, transform=transform_train, sgm=25, ratio=0.9, size_data=size_data, size_window=size_window)
         dataset_val = Dataset(dir_data_val, data_type=self.data_type, transform=transform_val, sgm=25, ratio=0.9, size_data=size_data, size_window=size_window)
 
-        loader_train = torch.utils.data.DataLoader(dataset_train, batch_size=batch_size, shuffle=True, num_workers=0)
-        loader_val = torch.utils.data.DataLoader(dataset_val, batch_size=batch_size, shuffle=True, num_workers=0)
+        loader_train = torch.utils.data.DataLoader(dataset_train, batch_size=batch_size, shuffle=True, num_workers=8)
+        loader_val = torch.utils.data.DataLoader(dataset_val, batch_size=batch_size, shuffle=True, num_workers=8)
 
         num_train = len(dataset_train)
         num_val = len(dataset_val)
@@ -356,7 +356,7 @@ def test(self):
 
         # dataset_test = Dataset(dir_data_test, data_type=self.data_type, transform=transform_test, sgm=(0, 25))
         dataset_test = Dataset(dir_data_test, data_type=self.data_type, transform=transform_test, sgm=25, ratio=1, size_data=size_data, size_window=size_window)
-        loader_test = torch.utils.data.DataLoader(dataset_test, batch_size=batch_size, shuffle=False, num_workers=0)
+        loader_test = torch.utils.data.DataLoader(dataset_test, batch_size=batch_size, shuffle=False, num_workers=8)
 
         num_test = len(dataset_test)