add debug

Author: scientist1642

main.py

@@ -15,6 +15,7 @@
 from test import test
 from utils import logger
 import my_optim
+import objgraph
 
 logger = logger.getLogger('main')
 
@@ -40,13 +41,16 @@
                     help='environment to train on (default: PongDeterministic-v3)')
 parser.add_argument('--no-shared', default=False, metavar='O',
                     help='use an optimizer without shared momentum.')
+parser.add_argument('--max-iters', type=int, default=50000,
+                    help='maximum iterations per process.')
 
+parser.add_argument('--debug', action='store_true', default=False,
+                    help='run in a way its easier to debug')
 
 if __name__ == '__main__':
     args = parser.parse_args()
 
     torch.manual_seed(args.seed)
-    #import ipdb; ipdb.set_trace()
     env = create_atari_env(args.env_name)
     shared_model = ActorCritic(
         env.observation_space.shape[0], env.action_space)
@@ -63,10 +67,15 @@
     p = mp.Process(target=test, args=(args.num_processes, args, shared_model))
     p.start()
     processes.append(p)
+    
+    if args.debug:
+        # run only one process in a main, easier to debug
+        train(0, args, shared_model, optimizer)
+    else:
+        for rank in range(0, args.num_processes):
+            p = mp.Process(target=train, args=(rank, args, shared_model, optimizer))
+            p.start()
+            processes.append(p)
 
-    for rank in range(0, args.num_processes):
-        p = mp.Process(target=train, args=(rank, args, shared_model, optimizer))
-        p.start()
-        processes.append(p)
     for p in processes:
         p.join()

train.py

@@ -11,6 +11,8 @@
 from torch.autograd import Variable
 from torchvision import datasets, transforms
 from utils import logger
+#import objgraph
+#from memory_profiler import profile
 
 
 logger = logger.getLogger('main')
@@ -22,7 +24,7 @@ def ensure_shared_grads(model, shared_model):
             return
         shared_param._grad = param.grad
 
-
+#@profile
 def train(rank, args, shared_model, optimizer=None):
     torch.manual_seed(args.seed + rank)
 
@@ -44,12 +46,18 @@ def train(rank, args, shared_model, optimizer=None):
 
     iteration = 0 
     while True:
+        if iteration == args.max_iters:
+            logger.info('Max iteration {} reached..'.format(args.max_iters))
+            break
 
-        iteration += 1  #TODO remove later
         if iteration % 100 == 0 and rank == 0:
-            logger.info('Memory usage of one proc: {} (mb)'.format(int(resource.getrusage(
-                resource.RUSAGE_SELF).ru_maxrss) / 1024 / 1024))
+            # for debugging purposes
+            mem_used = int(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss) 
+            mem_used_mb = mem_used / 1024 
 
+            logger.info('Memory usage of one proc: {} (mb)'.format(mem_used_mb))
+        
+        iteration += 1
         episode_length += 1
 
         # Sync with the shared model