add memory log

Author: scientist1642

main.py

@@ -13,8 +13,11 @@
 from model import ActorCritic
 from train import train
 from test import test
+from utils import logger
 import my_optim
 
+logger = logger.getLogger('main')
+
 # Based on
 # https://github.com/pytorch/examples/tree/master/mnist_hogwild
 # Training settings
@@ -43,7 +46,7 @@
     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)

test.py

@@ -11,7 +11,9 @@
 from torchvision import datasets, transforms
 import time
 from collections import deque
+from utils import logger
 
+logger = logger.getLogger('test')
 
 def test(rank, args, shared_model):
     torch.manual_seed(args.seed + rank)
@@ -59,7 +61,7 @@ def test(rank, args, shared_model):
             done = True
 
         if done:
-            print("Time {}, episode reward {}, episode length {}".format(
+            logger.info("Time {}, episode reward {}, episode length {}".format(
                 time.strftime("%Hh %Mm %Ss",
                               time.gmtime(time.time() - start_time)),
                 reward_sum, episode_length))

train.py

@@ -1,6 +1,7 @@
 import math
 import os
 import sys
+import resource 
 
 import torch
 import torch.nn.functional as F
@@ -9,6 +10,10 @@
 from model import ActorCritic
 from torch.autograd import Variable
 from torchvision import datasets, transforms
+from utils import logger
+
+
+logger = logger.getLogger('main')
 
 
 def ensure_shared_grads(model, shared_model):
@@ -36,8 +41,17 @@ def train(rank, args, shared_model, optimizer=None):
     done = True
 
     episode_length = 0
+
+    iteration = 0 
     while True:
+
+        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))
+
         episode_length += 1
+
         # Sync with the shared model
         model.load_state_dict(shared_model.state_dict())
         if done:

utils/logger.py

@@ -0,0 +1,49 @@
+# -*- coding: utf-8 -*-
+import os
+import logging
+import logging.config
+
+
+LOG_LEVEL = os.getenv('LOG_LEVEL', 'INFO')
+LOGGING = {
+    'version': 1,
+    'disable_existing_loggers': True,
+    'formatters': {
+        'verbose': {
+            'format': "[%(asctime)s] %(levelname)s " \
+                "[%(threadName)s:%(lineno)s] %(message)s",
+            'datefmt': "%Y-%m-%d %H:%M:%S"
+        },
+        'simple': {
+            'format': '%(levelname)s %(message)s'
+        },
+    },
+    'handlers': {
+        'console': {
+            'level': LOG_LEVEL,
+            'class': 'logging.StreamHandler',
+            'formatter': 'verbose'
+        },
+        'file': {
+            'level': LOG_LEVEL,
+            'class': 'logging.handlers.RotatingFileHandler',
+            'formatter': 'verbose',
+            'filename': 'rl.log',
+            'maxBytes': 10*10**6,
+            'backupCount': 3
+            }
+    },
+    'loggers': {
+        '': {
+            'handlers': ['console', 'file'],
+            'level': LOG_LEVEL,
+        },
+    }
+}
+
+
+logging.config.dictConfig(LOGGING)
+
+def getLogger(name):
+
+    return logging.getLogger(name)