We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
min_max log_gpu_memory option bug
**Describe the bug**
Setting `log_gpu_memory='min_max'` in `Trainer` leads to the following bug.

```
Traceback (most recent call last):
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 347, in fit
    self.single_gpu_train(model)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/dp_mixin.py", line 79, in single_gpu_train
    self.run_pretrain_routine(model)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 467, in run_pretrain_routine
    self.train()
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/train_loop_mixin.py", line 60, in train
    self.run_training_epoch()
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/train_loop_mixin.py", line 126, in run_training_epoch
    self.log_metrics(batch_step_metrics, grad_norm_dic)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/logging_mixin.py", line 20, in log_metrics
    mem_map = memory.get_memory_profile(self.log_gpu_memory)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/root_module/memory.py", line 205, in get_memory_profile
    for k, v in memory_map:
ValueError: too many values to unpack (expected 2)
```


**To Reproduce**
On current master, execute the following.
```
    trainer = Trainer(
        ...
        log_gpu_memory='min_max',
        ...
    )
    trainer.fit(model)
```

**Expected behavior**
Log the min/max utilization of gpu memory, as `min_max` option is documented.

**Desktop (please complete the following information):**
 - OS: Ubuntu 18.04 
 - Version: Current master

I am working on this issue. Will submit a PR soon.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `pytorch_lightning/root_module/memory.py`
Content:
```
1 '''
2 Generates a summary of a model's layers and dimensionality
3 '''
4 
5 import gc
6 import subprocess
7 
8 import numpy as np
9 import pandas as pd
10 import torch
11 
12 
13 class ModelSummary(object):
14 
15     def __init__(self, model, mode='full'):
16         '''
17         Generates summaries of model layers and dimensions.
18         '''
19         self.model = model
20         self.mode = mode
21         self.in_sizes = []
22         self.out_sizes = []
23 
24         self.summarize()
25 
26     def __str__(self):
27         return self.summary.__str__()
28 
29     def __repr__(self):
30         return self.summary.__str__()
31 
32     def named_modules(self):
33         if self.mode == 'full':
34             mods = self.model.named_modules()
35             mods = list(mods)[1:]  # do not include root module (LightningModule)
36         elif self.mode == 'top':
37             # the children are the top-level modules
38             mods = self.model.named_children()
39         else:
40             mods = []
41         return list(mods)
42 
43     def get_variable_sizes(self):
44         '''Run sample input through each layer to get output sizes'''
45         mods = self.named_modules()
46         in_sizes = []
47         out_sizes = []
48         input_ = self.model.example_input_array
49 
50         if self.model.on_gpu:
51             input_ = input_.cuda(0)
52 
53         if self.model.trainer.use_amp:
54             input_ = input_.half()
55 
56         with torch.no_grad():
57 
58             for _, m in mods:
59                 if type(input_) is list or type(input_) is tuple:  # pragma: no cover
60                     out = m(*input_)
61                 else:
62                     out = m(input_)
63 
64                 if type(input_) is tuple or type(input_) is list:  # pragma: no cover
65                     in_size = []
66                     for x in input_:
67                         if type(x) is list:
68                             in_size.append(len(x))
69                         else:
70                             in_size.append(x.size())
71                 else:
72                     in_size = np.array(input_.size())
73 
74                 in_sizes.append(in_size)
75 
76                 if type(out) is tuple or type(out) is list:  # pragma: no cover
77                     out_size = np.asarray([x.size() for x in out])
78                 else:
79                     out_size = np.array(out.size())
80 
81                 out_sizes.append(out_size)
82                 input_ = out
83 
84         self.in_sizes = in_sizes
85         self.out_sizes = out_sizes
86         assert len(in_sizes) == len(out_sizes)
87         return
88 
89     def get_layer_names(self):
90         '''Collect Layer Names'''
91         mods = self.named_modules()
92         names = []
93         layers = []
94         for name, m in mods:
95             names += [name]
96             layers += [str(m.__class__)]
97 
98         layer_types = [x.split('.')[-1][:-2] for x in layers]
99 
100         self.layer_names = names
101         self.layer_types = layer_types
102         return
103 
104     def get_parameter_sizes(self):
105         '''Get sizes of all parameters in `model`'''
106         mods = self.named_modules()
107         sizes = []
108         for _, m in mods:
109             p = list(m.parameters())
110             modsz = []
111             for j in range(len(p)):
112                 modsz.append(np.array(p[j].size()))
113             sizes.append(modsz)
114 
115         self.param_sizes = sizes
116         return
117 
118     def get_parameter_nums(self):
119         '''Get number of parameters in each layer'''
120         param_nums = []
121         for mod in self.param_sizes:
122             all_params = 0
123             for p in mod:
124                 all_params += np.prod(p)
125             param_nums.append(all_params)
126         self.param_nums = param_nums
127         return
128 
129     def make_summary(self):
130         '''
131         Makes a summary listing with:
132 
133         Layer Name, Layer Type, Input Size, Output Size, Number of Parameters
134         '''
135 
136         cols = ['Name', 'Type', 'Params']
137         if self.model.example_input_array is not None:
138             cols.extend(['In_sizes', 'Out_sizes'])
139 
140         df = pd.DataFrame(np.zeros((len(self.layer_names), len(cols))))
141         df.columns = cols
142 
143         df['Name'] = self.layer_names
144         df['Type'] = self.layer_types
145         df['Params'] = self.param_nums
146         df['Params'] = df['Params'].map(get_human_readable_count)
147 
148         if self.model.example_input_array is not None:
149             df['In_sizes'] = self.in_sizes
150             df['Out_sizes'] = self.out_sizes
151 
152         self.summary = df
153         return
154 
155     def summarize(self):
156         self.get_layer_names()
157         self.get_parameter_sizes()
158         self.get_parameter_nums()
159 
160         if self.model.example_input_array is not None:
161             self.get_variable_sizes()
162         self.make_summary()
163 
164 
165 def print_mem_stack():  # pragma: no cover
166     for obj in gc.get_objects():
167         try:
168             if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):
169                 print(type(obj), obj.size())
170         except Exception:
171             pass
172 
173 
174 def count_mem_items():  # pragma: no cover
175     nb_params = 0
176     nb_tensors = 0
177     for obj in gc.get_objects():
178         try:
179             if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):
180                 obj_type = str(type(obj))
181                 if 'parameter' in obj_type:
182                     nb_params += 1
183                 else:
184                     nb_tensors += 1
185         except Exception:
186             pass
187 
188     return nb_params, nb_tensors
189 
190 
191 def get_memory_profile(mode):
192     """
193     'all' means return memory for all gpus
194     'min_max' means return memory for max and min
195     :param mode:
196     :return:
197     """
198     memory_map = get_gpu_memory_map()
199 
200     if mode == 'min_max':
201         min_mem = 1000000
202         min_k = None
203         max_mem = 0
204         max_k = None
205         for k, v in memory_map:
206             if v > max_mem:
207                 max_mem = v
208                 max_k = k
209             if v < min_mem:
210                 min_mem = v
211                 min_k = k
212 
213         memory_map = {min_k: min_mem, max_k: max_mem}
214 
215     return memory_map
216 
217 
218 def get_gpu_memory_map():
219     """Get the current gpu usage.
220 
221     Returns
222     -------
223     usage: dict
224         Keys are device ids as integers.
225         Values are memory usage as integers in MB.
226     """
227     result = subprocess.check_output(
228         [
229             'nvidia-smi', '--query-gpu=memory.used',
230             '--format=csv,nounits,noheader'
231         ], encoding='utf-8')
232     # Convert lines into a dictionary
233     gpu_memory = [int(x) for x in result.strip().split('\n')]
234     gpu_memory_map = {}
235     for k, v in zip(range(len(gpu_memory)), gpu_memory):
236         k = f'gpu_{k}'
237         gpu_memory_map[k] = v
238     return gpu_memory_map
239 
240 
241 def get_human_readable_count(number):
242     """
243     Abbreviates an integer number with K, M, B, T for thousands, millions,
244     billions and trillions, respectively.
245     Examples:
246         123     -> 123
247         1234    -> 1 K       (one thousand)
248         2e6     -> 2 M       (two million)
249         3e9     -> 3 B       (three billion)
250         4e12    -> 4 T       (four trillion)
251         5e15    -> 5,000 T
252     :param number: a positive integer number
253     :returns a string formatted according to the pattern described above.
254     """
255     assert number >= 0
256     labels = [' ', 'K', 'M', 'B', 'T']
257     num_digits = int(np.floor(np.log10(number)) + 1 if number > 0 else 1)
258     num_groups = int(np.ceil(num_digits / 3))
259     num_groups = min(num_groups, len(labels))  # don't abbreviate beyond trillions
260     shift = -3 * (num_groups - 1)
261     number = number * (10 ** shift)
262     index = num_groups - 1
263     return f'{int(number):,d} {labels[index]}'
264
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/pytorch_lightning/root_module/memory.py b/pytorch_lightning/root_module/memory.py
--- a/pytorch_lightning/root_module/memory.py
+++ b/pytorch_lightning/root_module/memory.py
@@ -3,6 +3,7 @@
 '''
 
 import gc
+import os
 import subprocess
 
 import numpy as np
@@ -198,19 +199,10 @@
     memory_map = get_gpu_memory_map()
 
     if mode == 'min_max':
-        min_mem = 1000000
-        min_k = None
-        max_mem = 0
-        max_k = None
-        for k, v in memory_map:
-            if v > max_mem:
-                max_mem = v
-                max_k = k
-            if v < min_mem:
-                min_mem = v
-                min_k = k
-
-        memory_map = {min_k: min_mem, max_k: max_mem}
+        min_index, min_memory = min(memory_map.items(), key=lambda item: item[1])
+        max_index, max_memory = max(memory_map.items(), key=lambda item: item[1])
+
+        memory_map = {min_index: min_memory, max_index: max_memory}
 
     return memory_map
 
@@ -224,17 +216,18 @@
         Keys are device ids as integers.
         Values are memory usage as integers in MB.
     """
-    result = subprocess.check_output(
+    result = subprocess.run(
         [
-            'nvidia-smi', '--query-gpu=memory.used',
-            '--format=csv,nounits,noheader'
-        ], encoding='utf-8')
+            'nvidia-smi',
+            '--query-gpu=memory.used',
+            '--format=csv,nounits,noheader',
+        ],
+        encoding='utf-8',
+        capture_output=True,
+        check=True)
     # Convert lines into a dictionary
-    gpu_memory = [int(x) for x in result.strip().split('\n')]
-    gpu_memory_map = {}
-    for k, v in zip(range(len(gpu_memory)), gpu_memory):
-        k = f'gpu_{k}'
-        gpu_memory_map[k] = v
+    gpu_memory = [int(x) for x in result.stdout.strip().split(os.linesep)]
+    gpu_memory_map = {f'gpu_{index}': memory for index, memory in enumerate(gpu_memory)}
     return gpu_memory_map

{"golden_diff": "diff --git a/pytorch_lightning/root_module/memory.py b/pytorch_lightning/root_module/memory.py\n--- a/pytorch_lightning/root_module/memory.py\n+++ b/pytorch_lightning/root_module/memory.py\n@@ -3,6 +3,7 @@\n '''\n \n import gc\n+import os\n import subprocess\n \n import numpy as np\n@@ -198,19 +199,10 @@\n     memory_map = get_gpu_memory_map()\n \n     if mode == 'min_max':\n-        min_mem = 1000000\n-        min_k = None\n-        max_mem = 0\n-        max_k = None\n-        for k, v in memory_map:\n-            if v > max_mem:\n-                max_mem = v\n-                max_k = k\n-            if v < min_mem:\n-                min_mem = v\n-                min_k = k\n-\n-        memory_map = {min_k: min_mem, max_k: max_mem}\n+        min_index, min_memory = min(memory_map.items(), key=lambda item: item[1])\n+        max_index, max_memory = max(memory_map.items(), key=lambda item: item[1])\n+\n+        memory_map = {min_index: min_memory, max_index: max_memory}\n \n     return memory_map\n \n@@ -224,17 +216,18 @@\n         Keys are device ids as integers.\n         Values are memory usage as integers in MB.\n     \"\"\"\n-    result = subprocess.check_output(\n+    result = subprocess.run(\n         [\n-            'nvidia-smi', '--query-gpu=memory.used',\n-            '--format=csv,nounits,noheader'\n-        ], encoding='utf-8')\n+            'nvidia-smi',\n+            '--query-gpu=memory.used',\n+            '--format=csv,nounits,noheader',\n+        ],\n+        encoding='utf-8',\n+        capture_output=True,\n+        check=True)\n     # Convert lines into a dictionary\n-    gpu_memory = [int(x) for x in result.strip().split('\\n')]\n-    gpu_memory_map = {}\n-    for k, v in zip(range(len(gpu_memory)), gpu_memory):\n-        k = f'gpu_{k}'\n-        gpu_memory_map[k] = v\n+    gpu_memory = [int(x) for x in result.stdout.strip().split(os.linesep)]\n+    gpu_memory_map = {f'gpu_{index}': memory for index, memory in enumerate(gpu_memory)}\n     return gpu_memory_map\n", "issue": "min_max log_gpu_memory option bug\n**Describe the bug**\r\nSetting `log_gpu_memory='min_max'` in `Trainer` leads to the following bug.\r\n\r\n```\r\nTraceback (most recent call last):\r\n  File \"/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py\", line 347, in fit\r\n    self.single_gpu_train(model)\r\n  File \"/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/dp_mixin.py\", line 79, in single_gpu_train\r\n    self.run_pretrain_routine(model)\r\n  File \"/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py\", line 467, in run_pretrain_routine\r\n    self.train()\r\n  File \"/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/train_loop_mixin.py\", line 60, in train\r\n    self.run_training_epoch()\r\n  File \"/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/train_loop_mixin.py\", line 126, in run_training_epoch\r\n    self.log_metrics(batch_step_metrics, grad_norm_dic)\r\n  File \"/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/logging_mixin.py\", line 20, in log_metrics\r\n    mem_map = memory.get_memory_profile(self.log_gpu_memory)\r\n  File \"/opt/conda/lib/python3.7/site-packages/pytorch_lightning/root_module/memory.py\", line 205, in get_memory_profile\r\n    for k, v in memory_map:\r\nValueError: too many values to unpack (expected 2)\r\n```\r\n\r\n\r\n**To Reproduce**\r\nOn current master, execute the following.\r\n```\r\n    trainer = Trainer(\r\n        ...\r\n        log_gpu_memory='min_max',\r\n        ...\r\n    )\r\n    trainer.fit(model)\r\n```\r\n\r\n**Expected behavior**\r\nLog the min/max utilization of gpu memory, as `min_max` option is documented.\r\n\r\n**Desktop (please complete the following information):**\r\n - OS: Ubuntu 18.04 \r\n - Version: Current master\r\n\r\nI am working on this issue. Will submit a PR soon.\n", "before_files": [{"content": "'''\nGenerates a summary of a model's layers and dimensionality\n'''\n\nimport gc\nimport subprocess\n\nimport numpy as np\nimport pandas as pd\nimport torch\n\n\nclass ModelSummary(object):\n\n    def __init__(self, model, mode='full'):\n        '''\n        Generates summaries of model layers and dimensions.\n        '''\n        self.model = model\n        self.mode = mode\n        self.in_sizes = []\n        self.out_sizes = []\n\n        self.summarize()\n\n    def __str__(self):\n        return self.summary.__str__()\n\n    def __repr__(self):\n        return self.summary.__str__()\n\n    def named_modules(self):\n        if self.mode == 'full':\n            mods = self.model.named_modules()\n            mods = list(mods)[1:]  # do not include root module (LightningModule)\n        elif self.mode == 'top':\n            # the children are the top-level modules\n            mods = self.model.named_children()\n        else:\n            mods = []\n        return list(mods)\n\n    def get_variable_sizes(self):\n        '''Run sample input through each layer to get output sizes'''\n        mods = self.named_modules()\n        in_sizes = []\n        out_sizes = []\n        input_ = self.model.example_input_array\n\n        if self.model.on_gpu:\n            input_ = input_.cuda(0)\n\n        if self.model.trainer.use_amp:\n            input_ = input_.half()\n\n        with torch.no_grad():\n\n            for _, m in mods:\n                if type(input_) is list or type(input_) is tuple:  # pragma: no cover\n                    out = m(*input_)\n                else:\n                    out = m(input_)\n\n                if type(input_) is tuple or type(input_) is list:  # pragma: no cover\n                    in_size = []\n                    for x in input_:\n                        if type(x) is list:\n                            in_size.append(len(x))\n                        else:\n                            in_size.append(x.size())\n                else:\n                    in_size = np.array(input_.size())\n\n                in_sizes.append(in_size)\n\n                if type(out) is tuple or type(out) is list:  # pragma: no cover\n                    out_size = np.asarray([x.size() for x in out])\n                else:\n                    out_size = np.array(out.size())\n\n                out_sizes.append(out_size)\n                input_ = out\n\n        self.in_sizes = in_sizes\n        self.out_sizes = out_sizes\n        assert len(in_sizes) == len(out_sizes)\n        return\n\n    def get_layer_names(self):\n        '''Collect Layer Names'''\n        mods = self.named_modules()\n        names = []\n        layers = []\n        for name, m in mods:\n            names += [name]\n            layers += [str(m.__class__)]\n\n        layer_types = [x.split('.')[-1][:-2] for x in layers]\n\n        self.layer_names = names\n        self.layer_types = layer_types\n        return\n\n    def get_parameter_sizes(self):\n        '''Get sizes of all parameters in `model`'''\n        mods = self.named_modules()\n        sizes = []\n        for _, m in mods:\n            p = list(m.parameters())\n            modsz = []\n            for j in range(len(p)):\n                modsz.append(np.array(p[j].size()))\n            sizes.append(modsz)\n\n        self.param_sizes = sizes\n        return\n\n    def get_parameter_nums(self):\n        '''Get number of parameters in each layer'''\n        param_nums = []\n        for mod in self.param_sizes:\n            all_params = 0\n            for p in mod:\n                all_params += np.prod(p)\n            param_nums.append(all_params)\n        self.param_nums = param_nums\n        return\n\n    def make_summary(self):\n        '''\n        Makes a summary listing with:\n\n        Layer Name, Layer Type, Input Size, Output Size, Number of Parameters\n        '''\n\n        cols = ['Name', 'Type', 'Params']\n        if self.model.example_input_array is not None:\n            cols.extend(['In_sizes', 'Out_sizes'])\n\n        df = pd.DataFrame(np.zeros((len(self.layer_names), len(cols))))\n        df.columns = cols\n\n        df['Name'] = self.layer_names\n        df['Type'] = self.layer_types\n        df['Params'] = self.param_nums\n        df['Params'] = df['Params'].map(get_human_readable_count)\n\n        if self.model.example_input_array is not None:\n            df['In_sizes'] = self.in_sizes\n            df['Out_sizes'] = self.out_sizes\n\n        self.summary = df\n        return\n\n    def summarize(self):\n        self.get_layer_names()\n        self.get_parameter_sizes()\n        self.get_parameter_nums()\n\n        if self.model.example_input_array is not None:\n            self.get_variable_sizes()\n        self.make_summary()\n\n\ndef print_mem_stack():  # pragma: no cover\n    for obj in gc.get_objects():\n        try:\n            if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):\n                print(type(obj), obj.size())\n        except Exception:\n            pass\n\n\ndef count_mem_items():  # pragma: no cover\n    nb_params = 0\n    nb_tensors = 0\n    for obj in gc.get_objects():\n        try:\n            if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):\n                obj_type = str(type(obj))\n                if 'parameter' in obj_type:\n                    nb_params += 1\n                else:\n                    nb_tensors += 1\n        except Exception:\n            pass\n\n    return nb_params, nb_tensors\n\n\ndef get_memory_profile(mode):\n    \"\"\"\n    'all' means return memory for all gpus\n    'min_max' means return memory for max and min\n    :param mode:\n    :return:\n    \"\"\"\n    memory_map = get_gpu_memory_map()\n\n    if mode == 'min_max':\n        min_mem = 1000000\n        min_k = None\n        max_mem = 0\n        max_k = None\n        for k, v in memory_map:\n            if v > max_mem:\n                max_mem = v\n                max_k = k\n            if v < min_mem:\n                min_mem = v\n                min_k = k\n\n        memory_map = {min_k: min_mem, max_k: max_mem}\n\n    return memory_map\n\n\ndef get_gpu_memory_map():\n    \"\"\"Get the current gpu usage.\n\n    Returns\n    -------\n    usage: dict\n        Keys are device ids as integers.\n        Values are memory usage as integers in MB.\n    \"\"\"\n    result = subprocess.check_output(\n        [\n            'nvidia-smi', '--query-gpu=memory.used',\n            '--format=csv,nounits,noheader'\n        ], encoding='utf-8')\n    # Convert lines into a dictionary\n    gpu_memory = [int(x) for x in result.strip().split('\\n')]\n    gpu_memory_map = {}\n    for k, v in zip(range(len(gpu_memory)), gpu_memory):\n        k = f'gpu_{k}'\n        gpu_memory_map[k] = v\n    return gpu_memory_map\n\n\ndef get_human_readable_count(number):\n    \"\"\"\n    Abbreviates an integer number with K, M, B, T for thousands, millions,\n    billions and trillions, respectively.\n    Examples:\n        123     -> 123\n        1234    -> 1 K       (one thousand)\n        2e6     -> 2 M       (two million)\n        3e9     -> 3 B       (three billion)\n        4e12    -> 4 T       (four trillion)\n        5e15    -> 5,000 T\n    :param number: a positive integer number\n    :returns a string formatted according to the pattern described above.\n    \"\"\"\n    assert number >= 0\n    labels = [' ', 'K', 'M', 'B', 'T']\n    num_digits = int(np.floor(np.log10(number)) + 1 if number > 0 else 1)\n    num_groups = int(np.ceil(num_digits / 3))\n    num_groups = min(num_groups, len(labels))  # don't abbreviate beyond trillions\n    shift = -3 * (num_groups - 1)\n    number = number * (10 ** shift)\n    index = num_groups - 1\n    return f'{int(number):,d} {labels[index]}'\n", "path": "pytorch_lightning/root_module/memory.py"}], "after_files": [{"content": "'''\nGenerates a summary of a model's layers and dimensionality\n'''\n\nimport gc\nimport os\nimport subprocess\n\nimport numpy as np\nimport pandas as pd\nimport torch\n\n\nclass ModelSummary(object):\n\n    def __init__(self, model, mode='full'):\n        '''\n        Generates summaries of model layers and dimensions.\n        '''\n        self.model = model\n        self.mode = mode\n        self.in_sizes = []\n        self.out_sizes = []\n\n        self.summarize()\n\n    def __str__(self):\n        return self.summary.__str__()\n\n    def __repr__(self):\n        return self.summary.__str__()\n\n    def named_modules(self):\n        if self.mode == 'full':\n            mods = self.model.named_modules()\n            mods = list(mods)[1:]  # do not include root module (LightningModule)\n        elif self.mode == 'top':\n            # the children are the top-level modules\n            mods = self.model.named_children()\n        else:\n            mods = []\n        return list(mods)\n\n    def get_variable_sizes(self):\n        '''Run sample input through each layer to get output sizes'''\n        mods = self.named_modules()\n        in_sizes = []\n        out_sizes = []\n        input_ = self.model.example_input_array\n\n        if self.model.on_gpu:\n            input_ = input_.cuda(0)\n\n        if self.model.trainer.use_amp:\n            input_ = input_.half()\n\n        with torch.no_grad():\n\n            for _, m in mods:\n                if type(input_) is list or type(input_) is tuple:  # pragma: no cover\n                    out = m(*input_)\n                else:\n                    out = m(input_)\n\n                if type(input_) is tuple or type(input_) is list:  # pragma: no cover\n                    in_size = []\n                    for x in input_:\n                        if type(x) is list:\n                            in_size.append(len(x))\n                        else:\n                            in_size.append(x.size())\n                else:\n                    in_size = np.array(input_.size())\n\n                in_sizes.append(in_size)\n\n                if type(out) is tuple or type(out) is list:  # pragma: no cover\n                    out_size = np.asarray([x.size() for x in out])\n                else:\n                    out_size = np.array(out.size())\n\n                out_sizes.append(out_size)\n                input_ = out\n\n        self.in_sizes = in_sizes\n        self.out_sizes = out_sizes\n        assert len(in_sizes) == len(out_sizes)\n        return\n\n    def get_layer_names(self):\n        '''Collect Layer Names'''\n        mods = self.named_modules()\n        names = []\n        layers = []\n        for name, m in mods:\n            names += [name]\n            layers += [str(m.__class__)]\n\n        layer_types = [x.split('.')[-1][:-2] for x in layers]\n\n        self.layer_names = names\n        self.layer_types = layer_types\n        return\n\n    def get_parameter_sizes(self):\n        '''Get sizes of all parameters in `model`'''\n        mods = self.named_modules()\n        sizes = []\n        for _, m in mods:\n            p = list(m.parameters())\n            modsz = []\n            for j in range(len(p)):\n                modsz.append(np.array(p[j].size()))\n            sizes.append(modsz)\n\n        self.param_sizes = sizes\n        return\n\n    def get_parameter_nums(self):\n        '''Get number of parameters in each layer'''\n        param_nums = []\n        for mod in self.param_sizes:\n            all_params = 0\n            for p in mod:\n                all_params += np.prod(p)\n            param_nums.append(all_params)\n        self.param_nums = param_nums\n        return\n\n    def make_summary(self):\n        '''\n        Makes a summary listing with:\n\n        Layer Name, Layer Type, Input Size, Output Size, Number of Parameters\n        '''\n\n        cols = ['Name', 'Type', 'Params']\n        if self.model.example_input_array is not None:\n            cols.extend(['In_sizes', 'Out_sizes'])\n\n        df = pd.DataFrame(np.zeros((len(self.layer_names), len(cols))))\n        df.columns = cols\n\n        df['Name'] = self.layer_names\n        df['Type'] = self.layer_types\n        df['Params'] = self.param_nums\n        df['Params'] = df['Params'].map(get_human_readable_count)\n\n        if self.model.example_input_array is not None:\n            df['In_sizes'] = self.in_sizes\n            df['Out_sizes'] = self.out_sizes\n\n        self.summary = df\n        return\n\n    def summarize(self):\n        self.get_layer_names()\n        self.get_parameter_sizes()\n        self.get_parameter_nums()\n\n        if self.model.example_input_array is not None:\n            self.get_variable_sizes()\n        self.make_summary()\n\n\ndef print_mem_stack():  # pragma: no cover\n    for obj in gc.get_objects():\n        try:\n            if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):\n                print(type(obj), obj.size())\n        except Exception:\n            pass\n\n\ndef count_mem_items():  # pragma: no cover\n    nb_params = 0\n    nb_tensors = 0\n    for obj in gc.get_objects():\n        try:\n            if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):\n                obj_type = str(type(obj))\n                if 'parameter' in obj_type:\n                    nb_params += 1\n                else:\n                    nb_tensors += 1\n        except Exception:\n            pass\n\n    return nb_params, nb_tensors\n\n\ndef get_memory_profile(mode):\n    \"\"\"\n    'all' means return memory for all gpus\n    'min_max' means return memory for max and min\n    :param mode:\n    :return:\n    \"\"\"\n    memory_map = get_gpu_memory_map()\n\n    if mode == 'min_max':\n        min_index, min_memory = min(memory_map.items(), key=lambda item: item[1])\n        max_index, max_memory = max(memory_map.items(), key=lambda item: item[1])\n\n        memory_map = {min_index: min_memory, max_index: max_memory}\n\n    return memory_map\n\n\ndef get_gpu_memory_map():\n    \"\"\"Get the current gpu usage.\n\n    Returns\n    -------\n    usage: dict\n        Keys are device ids as integers.\n        Values are memory usage as integers in MB.\n    \"\"\"\n    result = subprocess.run(\n        [\n            'nvidia-smi',\n            '--query-gpu=memory.used',\n            '--format=csv,nounits,noheader',\n        ],\n        encoding='utf-8',\n        capture_output=True,\n        check=True)\n    # Convert lines into a dictionary\n    gpu_memory = [int(x) for x in result.stdout.strip().split(os.linesep)]\n    gpu_memory_map = {f'gpu_{index}': memory for index, memory in enumerate(gpu_memory)}\n    return gpu_memory_map\n\n\ndef get_human_readable_count(number):\n    \"\"\"\n    Abbreviates an integer number with K, M, B, T for thousands, millions,\n    billions and trillions, respectively.\n    Examples:\n        123     -> 123\n        1234    -> 1 K       (one thousand)\n        2e6     -> 2 M       (two million)\n        3e9     -> 3 B       (three billion)\n        4e12    -> 4 T       (four trillion)\n        5e15    -> 5,000 T\n    :param number: a positive integer number\n    :returns a string formatted according to the pattern described above.\n    \"\"\"\n    assert number >= 0\n    labels = [' ', 'K', 'M', 'B', 'T']\n    num_digits = int(np.floor(np.log10(number)) + 1 if number > 0 else 1)\n    num_groups = int(np.ceil(num_digits / 3))\n    num_groups = min(num_groups, len(labels))  # don't abbreviate beyond trillions\n    shift = -3 * (num_groups - 1)\n    number = number * (10 ** shift)\n    index = num_groups - 1\n    return f'{int(number):,d} {labels[index]}'\n", "path": "pytorch_lightning/root_module/memory.py"}]}