refactoring. encapsulate train()

RForestLiu · Aug 12, 2019 · e6b08d1 · e6b08d1
1 parent 63be7f4
commit e6b08d1
Showing 1 changed file with 19 additions and 22 deletions.
diff --git a/train.py b/train.py
@@ -22,25 +22,22 @@
 from lib.model import Ganomaly
 
 ##
-# def main():
-""" Training
-"""
-
-##
-# ARGUMENTS
-opt = Options().parse()
-
-##
-# LOAD DATA
-dataloader = load_data(opt)
-
-##
-# LOAD MODEL
-model = Ganomaly(opt, dataloader)
-
-##
-# TRAIN MODEL
-model.train()
-
-# if __name__ == '__main__':
-#     main()
+def train():
+    """ Training
+    """
+
+    ##
+    # ARGUMENTS
+    opt = Options().parse()
+    ##
+    # LOAD DATA
+    dataloader = load_data(opt)
+    ##
+    # LOAD MODEL
+    model = Ganomaly(opt, dataloader)
+    ##
+    # TRAIN MODEL
+    model.train()
+
+if __name__ == '__main__':
+    train()