Several improvement

train.py

@@ -1,3 +1,6 @@
+import os
+import resource
+import subprocess
 from typing import List, Tuple
 
 import torch
@@ -18,3 +21,15 @@ def parameter_summary(network: torch.nn.Module) -> List[Tuple[str, Tuple[int], s
         numpy = param.detach().cpu().numpy()
         parameter_info.append((name, numpy.shape, human_size(numpy.size * numpy.dtype.itemsize)))
     return parameter_info
+
+
+def resource_usage() -> Tuple[int, str]:
+    memory_peak = int(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss)
+    gpu_memory = subprocess.check_output(
+        'nvidia-smi --query-gpu=memory.used --format=csv,noheader', shell=True).decode()
+    if 'CUDA_VISIBLE_DEVICES' in os.environ:
+        gpu_memory = gpu_memory.split('\n')[int(os.environ['CUDA_VISIBLE_DEVICES'])]
+    else:
+        gpu_memory = ' '.join(gpu_memory.split('\n'))
+
+    return memory_peak, gpu_memory