Re-organize fgsa.py, naive_bayes.py

Moeed · Moeed · commit c89df6e8ab21 · 2019-09-25T12:42:41.000+02:00
Add torch to requirements.txt
diff --git a/fgsa.py b/fgsa.py
@@ -10,7 +10,7 @@
 
 epsilons = [0, .05, .1, .15, .2, .25, .3]
 use_cuda = True
-pretrained_model = "../data/lenet_mnist_model.pth"
+pretrained_model = "data/lenet_mnist_model.pth"
 
 
 # model under attack
diff --git a/naive_bayes.py b/naive_bayes.py
@@ -83,7 +83,7 @@ def test_model(model, priors, test):
 
 
 def main(train_size=.7):
-    train, test = preprocess_data(train_size, '../data/spambase.data')
+    train, test = preprocess_data(train_size, 'data/spambase.data')
     model = train_model(train)
     priors = get_priors(train)
     accuracy = test_model(model, priors, test)
diff --git a/requirements.txt b/requirements.txt
@@ -3,4 +3,6 @@ fire
 opencv-python
 matplotlib
 keras
-tensorflow-gpu
+tensorflow-gpu
+torch
+torchvision
