heojae
commented
3 years ago 원본 코드 변형 및 디렉토리 구조
저는 학습이 아닌 inference 관점에서, 분석하기 위해 원본 깃 레포 에서 필요한 부분만 빼내어서 변형하여 실험 하였습니다.
이걸 현재 레포에 추가할까 고민했지만, 레포와는 별개의 실험이니, 요기에 정리하기로 하였습니다.
- 폴더 구조
article_implem.py
utils
- memory.py
- plot.pyarticle_implem.py
접기/펼치기 버튼
```python
import pandas as pd
import torch
from torchvision.models import resnet18
from utils.memory import log_mem, getMemoryUsage
from utils.plot import plot_mem
# TODO : 조건을 선택해서, 실험해보세요.
base_dir = '.'
is_no_grad = False
model = resnet18().cuda()
model.eval()
# torch.save(model, "resnet18.pt") # 모델 용량 확인용, 46.8MB
print("------------------------------------ after upload model to gpu ---------------------------------------------")
print(f"Allocated: {round(torch.cuda.memory_allocated(0) / 1024 ** 2, 1)}MB, "
f"'Max Allocated: {round(torch.cuda.max_memory_allocated(0) / 1024 ** 2, 1)}MB, "
f" Cached: {round(torch.cuda.memory_cached(0) / 1024 ** 2, 1)}MB, "
f" Max Cached: {round(torch.cuda.max_memory_cached(0) / 1024 ** 2, 1)} MB, NVIDIA_SMI: {getMemoryUsage()} MB")
bs = 1
input = torch.rand(bs, 3, 224, 224).cuda()
mem_log = []
print("------------------------------------ after upload input to gpu ---------------------------------------------")
print(f"Allocated: {round(torch.cuda.memory_allocated(0) / 1024 ** 2, 1)}MB, "
f"'Max Allocated: {round(torch.cuda.max_memory_allocated(0) / 1024 ** 2, 1)}MB, "
f" Cached: {round(torch.cuda.memory_cached(0) / 1024 ** 2, 1)}MB, "
f" Max Cached: {round(torch.cuda.max_memory_cached(0) / 1024 ** 2, 1)} MB, NVIDIA_SMI: {getMemoryUsage()} MB")
try:
if is_no_grad:
with torch.no_grad():
no_grad = "_no_grad"
print("------------------------------------ 1 Epoch ----------------------------------")
mem_log.extend(log_mem(model, input, exp='1'))
print("------------------------------------ 2 Epoch ----------------------------------")
mem_log.extend(log_mem(model, input, exp='2'))
print("------------------------------------ 3 Epoch ----------------------------------")
mem_log.extend(log_mem(model, input, exp='3'))
else:
no_grad = ""
print("------------------------------------ 1 Epoch ----------------------------------")
mem_log.extend(log_mem(model, input, exp='1'))
print("------------------------------------ 2 Epoch ----------------------------------")
mem_log.extend(log_mem(model, input, exp='2'))
print("------------------------------------ 3 Epoch ----------------------------------")
mem_log.extend(log_mem(model, input, exp='3'))
except Exception as e:
print(e)
print(f'log_mem failed because of {e}')
df = pd.DataFrame(mem_log)
plot_mem(df, exps=['1', '2', '3'], output_file=f'{base_dir}/baseline_memory_plot_{bs}{no_grad}.png')
```
plot.py
접기/펼치기 버튼
```python
from matplotlib import pyplot as plt
def plot_mem(
df,
exps=None,
normalize_call_idx=True,
normalize_mem_all=True,
filter_fwd=False,
return_df=False,
output_file=None
):
if exps is None:
exps = df.exp.drop_duplicates()
fig, ax = plt.subplots(figsize=(20, 10))
for exp in exps:
df_ = df[df.exp == exp]
if normalize_call_idx:
df_.call_idx = df_.call_idx / df_.call_idx.max()
if normalize_mem_all:
df_.mem_all = df_.mem_all - df_[df_.call_idx == df_.call_idx.min()].mem_all.iloc[0]
df_.mem_all = df_.mem_all // 2 ** 20
if filter_fwd:
layer_idx = 0
callidx_stop = df_[(df_["layer_idx"] == layer_idx) & (df_["hook_type"] == "fwd")]["call_idx"].iloc[0]
df_ = df_[df_["call_idx"] <= callidx_stop]
# df_ = df_[df_.call_idx < df_[df_.layer_idx=='bwd'].call_idx.min()]
plot = df_.plot(ax=ax, x='call_idx', y='mem_all', label=exp)
if output_file:
plot.get_figure().savefig(output_file)
if return_df:
return df_
```
memory.py
접기/펼치기 버튼
```python
from torch.utils.checkpoint import checkpoint_sequential
import torch
from pynvml.smi import nvidia_smi
nvsmi = nvidia_smi.getInstance()
def getMemoryUsage():
usage = nvsmi.DeviceQuery("memory.used")["gpu"][0]["fb_memory_usage"]
return "%d %s" % (usage["used"], usage["unit"])
def _get_gpu_mem(synchronize=True, empty_cache=True):
print(f"Allocated: {round(torch.cuda.memory_allocated(0) / 1024 ** 2, 1)}MB, "
f"'Max Allocated: {round(torch.cuda.max_memory_allocated(0) / 1024 ** 2, 1)}MB, "
f" Cached: {round(torch.cuda.memory_cached(0) / 1024 ** 2, 1)}MB, "
f" Max Cached: {round(torch.cuda.max_memory_cached(0) / 1024 ** 2, 1)} MB, NVIDIA_SMI: {getMemoryUsage()} MB")
return torch.cuda.memory_allocated(), torch.cuda.memory_cached()
def _generate_mem_hook(handle_ref, mem, idx, hook_type, exp):
def hook(self, *args):
if len(mem) == 0 or mem[-1]["exp"] != exp:
call_idx = 0
else:
call_idx = mem[-1]["call_idx"] + 1
mem_all, mem_cached = _get_gpu_mem()
torch.cuda.synchronize()
mem.append({
'layer_idx': idx,
'call_idx': call_idx,
'layer_type': type(self).__name__,
'exp': exp,
'hook_type': hook_type,
'mem_all': mem_all,
'mem_cached': mem_cached,
})
return hook
def _add_memory_hooks(idx, mod, mem_log, exp, hr):
h = mod.register_forward_pre_hook(_generate_mem_hook(hr, mem_log, idx, 'pre', exp))
hr.append(h)
h = mod.register_forward_hook(_generate_mem_hook(hr, mem_log, idx, 'fwd', exp))
hr.append(h)
h = mod.register_backward_hook(_generate_mem_hook(hr, mem_log, idx, 'bwd', exp))
hr.append(h)
def log_mem(model, inp, mem_log=None, exp=None):
mem_log = mem_log or []
exp = exp or f'exp_{len(mem_log)}'
hr = []
for idx, module in enumerate(model.modules()):
_add_memory_hooks(idx, module, mem_log, exp, hr)
try:
out = model(inp)
# loss = out.sum()
# loss.backward()
finally:
[h.remove() for h in hr]
return mem_log
```
참고자료
https://www.sicara.ai/blog/2019-28-10-deep-learning-memory-usage-and-pytorch-optimization-tricks
위 글을 보면서,
학습을 할 때의GPU상에memory가 쌓이는 순서 및release되는 방향에 대해서, 수치적으로 이해를 할 수 있었습니다.하지만, 실제로
torch에서는cache에 대해서도 고려를 해주어야할 필요도 있고,with torch.no_grad()를 통해서,inference를 할 때 또한, 생각을 해주어야 할 필요가 있습니다.https://pytorch.org/docs/stable/notes/cuda.html#cuda-memory-management
그렇기에, 이를 참고하여, 실제로
train과inference를 할 때,GPU memory와torch memory에서는 어떠하게 다르며 이를 분석하고자, 아래와 같은 실험을 하게되었습니다. 대부분의model infernceAPI는with torch.no_grad()를 사용을 많이하게 되어있는 데, 이를 수치적으로 어떠한 만큼 효율을 낼 수 있는지 보여주고 싶습니다.