corentin/torch_utils
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Several improvement

Browse source
This commit is contained in:
Corentin Risselin 2020-04-13 17:30:35 +09:00
commit 12afc7cc93
2 changed files with 16 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

2
layers.py
View file

@ -48,7 +48,7 @@ class Layer(nn.Module):
class Conv2d(Layer): class Conv2d(Layer):
def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3, stride: int = 1,
activation=0, batch_norm=None, **kwargs): activation=0, batch_norm=None, **kwargs):
super(Conv2d, self).__init__(activation, batch_norm) super(Conv2d, self).__init__(activation, batch_norm)

15
train.py
View file

@ -1,3 +1,6 @@
import os
import resource
import subprocess
from typing import List, Tuple from typing import List, Tuple
import torch import torch
@ -18,3 +21,15 @@ def parameter_summary(network: torch.nn.Module) -> List[Tuple[str, Tuple[int], s
numpy = param.detach().cpu().numpy() numpy = param.detach().cpu().numpy()
parameter_info.append((name, numpy.shape, human_size(numpy.size * numpy.dtype.itemsize))) parameter_info.append((name, numpy.shape, human_size(numpy.size * numpy.dtype.itemsize)))
return parameter_info 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