Hi All,
I was experiencing massive memory usage on my Jetson Nano when running Pytorch 1.8 (installed from the first page of this forum), and I was wondering if anyone had any similar issues, or if there was a way around it?
I have a small script that I was profiling with memory_profiler
, it looks like this:
import torch
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('size', type=int)
parser.add_argument('--cpu', action='store_true')
args = parser.parse_args()
@profile
def f():
torch.set_grad_enabled(False)
torch.cuda._lazy_init()
device = 'cuda' if torch.cuda.is_available() else 'cpu'
if args.cpu:
device = 'cpu'
model = torch.nn.Conv2d(1, 1, 1).to(device)
x = torch.rand(1, 1, args.size, args.size).to(device)
y = model(x)
if __name__ == '__main__':
f()
When running the profiler without the cpu
flag (ie, using CUDA) with a size of 100, I get:
python3 -m memory_profiler torchmemscript.py 100
Filename: torchmemscript.py
Line # Mem usage Increment Occurences Line Contents
============================================================
9 150.906 MiB 150.906 MiB 1 @profile
10 def f():
11 150.906 MiB 0.000 MiB 1 torch.set_grad_enabled(False)
12 155.336 MiB 4.430 MiB 1 torch.cuda._lazy_init()
13 155.336 MiB 0.000 MiB 1 device = 'cuda' if torch.cuda.is_available() else 'cpu'
14 155.336 MiB 0.000 MiB 1 if args.cpu:
15 device = 'cpu'
16 1889.699 MiB 1734.363 MiB 1 model = torch.nn.Conv2d(1, 1, 1).to(device)
17 1890.414 MiB 0.715 MiB 1 x = torch.rand(1, 1, args.size, args.size).to(device)
18 2634.496 MiB 744.082 MiB 1 y = model(x)
But when running with the cpu flag I get:
python3 -m memory_profiler torchmemscript.py 100 --cpu
Filename: torchmemscript.py
Line # Mem usage Increment Occurences Line Contents
============================================================
9 151.055 MiB 151.055 MiB 1 @profile
10 def f():
11 151.055 MiB 0.000 MiB 1 torch.set_grad_enabled(False)
12 155.359 MiB 4.305 MiB 1 torch.cuda._lazy_init()
13 155.359 MiB 0.000 MiB 1 device = 'cuda' if torch.cuda.is_available() else 'cpu'
14 155.359 MiB 0.000 MiB 1 if args.cpu:
15 155.359 MiB 0.000 MiB 1 device = 'cpu'
16 157.754 MiB 2.395 MiB 1 model = torch.nn.Conv2d(1, 1, 1).to(device)
17 157.754 MiB 0.000 MiB 1 x = torch.rand(1, 1, args.size, args.size).to(device)
18 160.051 MiB 2.297 MiB 1 y = model(x)
Does anyone know why there is such a huge memory usage? It’s killing my torch apps, which otherwise don’t use much RAM at all