Merge pull request #6 from samuel920119920119samuel/feature/passing-visdom-arguments

feature: add vis_args for passing visdom arguments

Author: srush

project/datasets.py

@@ -22,10 +22,10 @@ def make_pts(N):
 
 
 class Graph:
-    def __init__(self, vis=False):
+    def __init__(self, vis=False, vis_args={}):
         self.gifs = []
         if vis:
-            self.vis = visdom.Visdom()
+            self.vis = visdom.Visdom(**vis_args)
         else:
             self.vis = None
         self.first = True
@@ -74,8 +74,8 @@ def graph(self, outfile, model=None):
 
 
 class Simple(Graph):
-    def __init__(self, N, vis=False):
-        super().__init__(vis)
+    def __init__(self, N, vis=False, vis_args={}):
+        super().__init__(vis, vis_args)
         self.N = N
         self.X = make_pts(N)
         self.y = []
@@ -85,8 +85,8 @@ def __init__(self, N, vis=False):
 
 
 class Split(Graph):
-    def __init__(self, N, vis=False):
-        super().__init__(vis)
+    def __init__(self, N, vis=False, vis_args={}):
+        super().__init__(vis, vis_args)
         self.N = N
         self.X = make_pts(N)
         self.y = []
@@ -96,8 +96,8 @@ def __init__(self, N, vis=False):
 
 
 class Xor(Graph):
-    def __init__(self, N, vis=False):
-        super().__init__(vis)
+    def __init__(self, N, vis=False, vis_args={}):
+        super().__init__(vis, vis_args)
         self.N = N
         self.X = make_pts(N)
         self.y = []