clean code

Author: AndyDeng

train_pytorch.py

@@ -32,6 +32,55 @@
 random.seed(0)
 dtype = torch.cuda.FloatTensor
 
+
+
+# Load Hyperparameters
+parser = argparse.ArgumentParser()
+parser.add_argument('--gpu', type=int, default=0, help='GPU to use [default: GPU 0]')
+parser.add_argument('--model', default='pointnet_cls',
+                    help='Model name: pointnet_cls or pointnet_cls_basic [default: pointnet_cls]')
+parser.add_argument('--log_dir', default='log', help='Log dir [default: log]')
+parser.add_argument('--num_point', type=int, default=1024, help='Point Number [256/512/1024/2048] [default: 1024]')
+parser.add_argument('--max_epoch', type=int, default=2, help='Epoch to run [default: 250]')
+parser.add_argument('--batch_size', type=int, default=32, help='Batch Size during training [default: 32]')
+parser.add_argument('--learning_rate', type=float, default=0.001, help='Initial learning rate [default: 0.001]')
+parser.add_argument('--momentum', type=float, default=0.9, help='Initial learning rate [default: 0.9]')
+parser.add_argument('--optimizer', default='adam', help='adam or momentum [default: adam]')
+parser.add_argument('--decay_step', type=int, default=200000, help='Decay step for lr decay [default: 200000]')
+parser.add_argument('--decay_rate', type=float, default=0.7, help='Decay rate for lr decay [default: 0.8]')
+FLAGS = parser.parse_args()
+
+NUM_POINT = FLAGS.num_point
+LEARNING_RATE = FLAGS.learning_rate
+GPU_INDEX = FLAGS.gpu
+MOMENTUM = FLAGS.momentum
+       
+MAX_NUM_POINT = 2048
+
+BN_INIT_DECAY = 0.5
+BN_DECAY_DECAY_RATE = 0.5
+#BN_DECAY_DECAY_STEP = float(DECAY_STEP)
+BN_DECAY_CLIP = 0.99
+
+decay_steps = FLAGS.decay_step
+decay_rate = FLAGS.decay_rate
+LEARNING_RATE_MIN = 0.00001
+        
+NUM_CLASS = 40
+#sample_num = 160
+BATCH_SIZE = FLAGS.batch_size #32
+NUM_EPOCHS = FLAGS.max_epoch
+jitter = 0.01
+jitter_val = 0.01
+
+rotation_range = [0, math.pi / 18, 0, 'g']
+rotation_rage_val = [0, 0, 0, 'u']
+order = 'rxyz'
+
+scaling_range = [0.05, 0.05, 0.05, 'g']
+scaling_range_val = [0, 0, 0, 'u']
+
+
 class modelnet40_dataset(Dataset):
 
     def __init__(self, data, labels):
@@ -44,30 +93,32 @@ def __len__(self):
     def __getitem__(self, i):
         return self.data[i], self.labels[i]
 
+
 # C_in, C_out, D, N_neighbors, dilution, N_rep, r_indices_func, C_lifted = None, mlp_width = 2
 # (a, b, c, d, e) == (C_in, C_out, N_neighbors, dilution, N_rep)
 # Abbreviated PointCNN constructor.
-AbbPointCNN = lambda a,b,c,d,e: RandPointCNN(a, b, 3, c, d, e, knn_indices_func_gpu)
+AbbPointCNN = lambda a, b, c, d, e: RandPointCNN(a, b, 3, c, d, e, knn_indices_func_gpu)
+
 
 class Classifier(nn.Module):
 
     def __init__(self):
         super(Classifier, self).__init__()
-        
-        self.pcnn1 = AbbPointCNN(  3,  32,  8, 1,  -1)
+
+        self.pcnn1 = AbbPointCNN(3, 32, 8, 1, -1)
         self.pcnn2 = nn.Sequential(
-            AbbPointCNN( 32,  64,  8, 2,  -1),
-            AbbPointCNN( 64,  96,  8, 4,  -1),
-            AbbPointCNN( 96, 128, 12, 4, 120),
+            AbbPointCNN(32, 64, 8, 2, -1),
+            AbbPointCNN(64, 96, 8, 4, -1),
+            AbbPointCNN(96, 128, 12, 4, 120),
             AbbPointCNN(128, 160, 12, 6, 120)
         )
-        
+
         self.fcn = nn.Sequential(
             Dense(160, 128),
-            Dense(128,  64, drop_rate = 0.5),
-            Dense( 64,  40, with_bn = False, activation = None)
+            Dense(128, 64, drop_rate=0.5),
+            Dense(64, NUM_CLASS, with_bn=False, activation=None)
         )
-        
+
     def forward(self, x):
         x = self.pcnn1(x)
         if False:
@@ -82,59 +133,14 @@ def forward(self, x):
         x = self.pcnn2(x)[1]  # grab features
 
         logits = self.fcn(x)
-        logits_mean = torch.mean(logits, dim = 1)
+        logits_mean = torch.mean(logits, dim=1)
         return logits_mean
 
-# Load Hyperparameters
-parser = argparse.ArgumentParser()
-parser.add_argument('--gpu', type=int, default=0, help='GPU to use [default: GPU 0]')
-parser.add_argument('--model', default='pointnet_cls',
-                    help='Model name: pointnet_cls or pointnet_cls_basic [default: pointnet_cls]')
-parser.add_argument('--log_dir', default='log', help='Log dir [default: log]')
-parser.add_argument('--num_point', type=int, default=1024, help='Point Number [256/512/1024/2048] [default: 1024]')
-parser.add_argument('--max_epoch', type=int, default=2, help='Epoch to run [default: 250]')
-parser.add_argument('--batch_size', type=int, default=32, help='Batch Size during training [default: 32]')
-parser.add_argument('--learning_rate', type=float, default=0.001, help='Initial learning rate [default: 0.001]')
-parser.add_argument('--momentum', type=float, default=0.9, help='Initial learning rate [default: 0.9]')
-parser.add_argument('--optimizer', default='adam', help='adam or momentum [default: adam]')
-parser.add_argument('--decay_step', type=int, default=200000, help='Decay step for lr decay [default: 200000]')
-parser.add_argument('--decay_rate', type=float, default=0.7, help='Decay rate for lr decay [default: 0.8]')
-FLAGS = parser.parse_args()
-
-NUM_POINT = FLAGS.num_point
-lr = FLAGS.learning_rate 
-GPU_INDEX = FLAGS.gpu
-MOMENTUM = FLAGS.momentum
-       
-MAX_NUM_POINT = 2048
-
-BN_INIT_DECAY = 0.5
-BN_DECAY_DECAY_RATE = 0.5
-#BN_DECAY_DECAY_STEP = float(DECAY_STEP)
-BN_DECAY_CLIP = 0.99        
-        
-        
-num_class = 40
-#sample_num = 160
-BATCH_SIZE = FLAGS.batch_size #32
-num_epochs = FLAGS.max_epoch 
-jitter = 0.01
-jitter_val = 0.01
-
-rotation_range = [0, math.pi / 18, 0, 'g']
-rotation_rage_val = [0, 0, 0, 'u']
-order = 'rxyz'
-
-scaling_range = [0.05, 0.05, 0.05, 'g']
-scaling_range_val = [0, 0, 0, 'u']
 
 print("------Building model-------")
 model = Classifier().cuda()
 print("------Successfully Built model-------")
 
-decay_steps = FLAGS.decay_step
-decay_rate = FLAGS.decay_rate  
-lr_min = 0.00001
 
 optimizer = torch.optim.SGD(model.parameters(), lr = 0.01, momentum = 0.9)
 loss_fn = nn.CrossEntropyLoss()
@@ -144,19 +150,19 @@ def forward(self, x):
 #model_save_dir = os.path.join(CURRENT_DIR, "models", "mnist2")
 #os.makedirs(model_save_dir, exist_ok = True)
 
-TRAIN_FILES = provider.getDataFiles( \
-    os.path.join(BASE_DIR, 'data/modelnet40_ply_hdf5_2048/train_files.txt'))
-TEST_FILES = provider.getDataFiles(\
-    os.path.join(BASE_DIR, 'data/modelnet40_ply_hdf5_2048/test_files.txt'))
+TRAIN_FILES = provider.getDataFiles(os.path.join(BASE_DIR, 'data/modelnet40_ply_hdf5_2048/train_files.txt'))
+TEST_FILES = provider.getDataFiles(os.path.join(BASE_DIR, 'data/modelnet40_ply_hdf5_2048/test_files.txt'))
 
 losses = []
 accuracies = []
 
+'''
 if False:
     latest_model = sorted(os.listdir(model_save_dir))[-1]
     model.load_state_dict(torch.load(os.path.join(model_save_dir, latest_model)))    
-     
-for epoch in range(1, num_epochs+1):   
+'''
+
+for epoch in range(1, NUM_EPOCHS+1):
     train_file_idxs = np.arange(0, len(TRAIN_FILES))
     np.random.shuffle(train_file_idxs)
 
@@ -176,10 +182,10 @@ def forward(self, x):
         loss_sum = 0
 
         if epoch > 1:
-            lr *= decay_rate ** (global_step // decay_steps)
-            if lr > lr_min:
-                print("NEW LEARNING RATE:", lr)
-                optimizer = torch.optim.SGD(model.parameters(), lr = lr, momentum = 0.9)
+            LEARNING_RATE *= decay_rate ** (global_step // decay_steps)
+            if LEARNING_RATE > LEARNING_RATE_MIN:
+                print("NEW LEARNING RATE:", LEARNING_RATE)
+                optimizer = torch.optim.SGD(model.parameters(), lr = LEARNING_RATE, momentum = 0.9)
 
         for batch_idx in range(num_batches):
             start_idx = batch_idx * BATCH_SIZE