normal_conv_mem.py

from typing import Any
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import autograd


def scientific_precision(number):
    suffix = ["KB", "MB", "GB", "TB", "PB"]
    for idx_, s in enumerate(suffix):
        idx = idx_ + 1
        if 1024 ** (idx + 1) < number:
            continue
        return number / 1024 ** idx, s
    return number / 1024 ** idx, s


def print_gpu_utilization(idx=None, prefix=None):
    from pynvml import nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo, nvmlInit

    nvmlInit()
    used = 0
    if idx is None:
        # sum all GPU memory usage
        n = torch.cuda.device_count()
        for index in range(n):
            handle = nvmlDeviceGetHandleByIndex(index)
            info = nvmlDeviceGetMemoryInfo(handle)
            used += info.used
    else:
        # report specific GPU memory usage only
        handle = nvmlDeviceGetHandleByIndex(index)
        info = nvmlDeviceGetMemoryInfo(handle)
        used += info.used
    sci_used, suff = scientific_precision(used)
    if prefix:
        print(f"[{prefix}] GPU memory occupied: {sci_used:.3f} {suff}.")
    else:
        print(f"GPU memory occupied: {sci_used:.3f} {suff}.")
    return used


net = nn.Sequential(*[nn.Conv2d(32, 32, kernel_size=1, bias=True),] * 50)
print("normal nn.conv2d")
print_gpu_utilization(prefix="Torch init")

device = "cuda"
data = torch.randn(16, 32, 224, 224).to(device).requires_grad_(True)
net = net.to(device)

for n, p in net.named_parameters():
    p.requires_grad_(False)


with torch.no_grad():
    out = net(data)
    print_gpu_utilization(prefix="Forward")

out = net(data)
out.sum().backward()
print_gpu_utilization(prefix="Training")