Fix CLI flags and seed handling

README.md

@@ -31,10 +31,10 @@ PyTorch implements "DeepLog: Anomaly Detection and Diagnosis from System Logs th
 
    2.2. Train
 
-   `num-class` is count of `event_id_map`, where `event_id_map` is generated by `preprocess.py`. `num-candidates` is self-define, here we define `num-candidates` is `num-class*0.1`
+   `num-classes` is count of `event_id_map`, where `event_id_map` is generated by `preprocess.py`. `num-candidates` is self-define, here we define `num-candidates` is `num-classes*0.1`
 
    ```python
-   python3 train.py --num-class 1143 --num-candidates 114 --epochs 35 --window-size 3 --local True
+   python3 train.py --num-classes 1143 --num-candidates 114 --epochs 35 --window-size 3 --local
    ```
 
    2.3. Predict

deeplog/deeplog.py

@@ -3,8 +3,9 @@
 import json
 import logging
 import argparse
+import random
 
-import boto3
+import numpy as np
 import torch
 import torch.distributed as dist
 import torch.nn as nn
@@ -38,6 +39,14 @@ def forward(self, input):
         return out
 
 
+def set_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
 class Generate():
     def __init__(self):
         self.init_obj = None
@@ -70,6 +79,7 @@ def init_line(self, local, name):
             self.init_obj = f
             line = self.init_obj.readline()
         else:
+            import boto3
             client = boto3.client('s3')
             bucket = BUCKET
             prefix = PREFIX
@@ -127,10 +137,9 @@ def train(args):
             dist.get_rank(), args.num_gpus))
 
     # set the seed for generating random numbers
-    torch.manual_seed(args.seed)
+    set_seed(args.seed)
     if use_cuda:
         logger.info('Use CUDA')
-        torch.cuda.manual_seed(args.seed)
 
     train_loader = _get_train_data_loader(args.batch_size, is_distributed, args.window_size, args.local, **kwargs)
 
@@ -300,7 +309,7 @@ def output_fn(prediction, accept):
     parser.add_argument('--num-gpus', type=int, default=os.environ['SM_NUM_GPUS'])
 
     # Local mode
-    parser.add_argument('--local', type=bool, default=False,
-                        help='local training model.')
+    parser.add_argument('--local', action='store_true',
+                        help='use local training files instead of S3.')
 
     train(parser.parse_args())

example/train.py

@@ -2,9 +2,10 @@
 import sys
 import logging
 import argparse
-sys.path.append('../')
-import torch
-torch.manual_seed(0)
+
+PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, PROJECT_ROOT)
+
 from deeplog.deeplog import train
 
 logging.basicConfig(level=logging.DEBUG,
@@ -49,8 +50,8 @@
                         help='number of gpu to train')
 
     # Local mode
-    parser.add_argument('--local', type=bool, default=False,
-                        help='local training model.')
+    parser.add_argument('--local', action='store_true',
+                        help='use local training files instead of S3.')
 
     if not os.path.isdir('./model/'):
         os.mkdir('./model/')