tiny refactor

Author: scientist1642

main.py

@@ -15,7 +15,6 @@
 from test import test
 from utils import logger
 import my_optim
-import objgraph
 
 logger = logger.getLogger('main')
 
@@ -62,20 +61,19 @@
         optimizer = my_optim.SharedAdam(shared_model.parameters(), lr=args.lr)
         optimizer.share_memory()
 
-    processes = []
-
-    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:
+    if not args.debug:
+        processes = []
+
+        p = mp.Process(target=test, args=(args.num_processes, args, shared_model))
+        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()
+        for p in processes:
+            p.join()
+    else: ## debug is enabled
+        # run only one process in a main, easier to debug
+        train(0, args, shared_model, optimizer)

model.py

@@ -44,8 +44,11 @@ def __init__(self, num_inputs, action_space):
         self.lstm = nn.LSTMCell(32 * 3 * 3, 256)
 
         num_outputs = action_space.n
+
         self.critic_linear = nn.Linear(256, 1)
         self.actor_linear = nn.Linear(256, num_outputs)
+        #self.critic_linear = nn.Linear(288, 1)
+        #self.actor_linear = nn.Linear(288, num_outputs)
 
         self.apply(weights_init)
         self.actor_linear.weight.data = normalized_columns_initializer(
@@ -66,7 +69,6 @@ def forward(self, inputs):
         x = F.elu(self.conv2(x))
         x = F.elu(self.conv3(x))
         x = F.elu(self.conv4(x))
-
         x = x.view(-1, 32 * 3 * 3)
         hx, cx = self.lstm(x, (hx, cx))
         x = hx

train.py

@@ -2,6 +2,7 @@
 import os
 import sys
 import resource 
+import gc
 
 import torch
 import torch.nn.functional as F
@@ -11,20 +12,15 @@
 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')
 
-
 def ensure_shared_grads(model, shared_model):
     for param, shared_param in zip(model.parameters(), shared_model.parameters()):
         if shared_param.grad is not None:
             return
         shared_param._grad = param.grad
 
-#@profile
 def train(rank, args, shared_model, optimizer=None):
     torch.manual_seed(args.seed + rank)
 
@@ -45,18 +41,24 @@ def train(rank, args, shared_model, optimizer=None):
     episode_length = 0
 
     iteration = 0 
+    
     while True:
+
+        values = []
+        log_probs = []
+        rewards = []
+        entropies = []
+        
         if iteration == args.max_iters:
             logger.info('Max iteration {} reached..'.format(args.max_iters))
             break
 
-        if iteration % 100 == 0 and rank == 0:
-            # for debugging purposes
+        if iteration % 200 == 0 and rank == 0:
             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
 
@@ -69,10 +71,6 @@ def train(rank, args, shared_model, optimizer=None):
             cx = Variable(cx.data)
             hx = Variable(hx.data)
 
-        values = []
-        log_probs = []
-        rewards = []
-        entropies = []
 
         for step in range(args.num_steps):
             value, logit, (hx, cx) = model(