pytorch - 如何减少 Pytorch 在使用 CUDA 时使用的(非 GPU)RAM 开销?
问题描述
我正在尝试在我的 Jetson Nano(4GB RAM)上运行一些 Pytorch 模型,但我了解到 Pytorch 使用大约 2GB 的 RAM 来初始化任何与 CUDA 相关的东西。
我已经做了一些测试(在这个 github 问题的帮助下),并运行了以下脚本:
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()
可以使用python3 -m memory_profiler torchmemscript.py 100
. 这是输出:
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)
很明显,模型已加载并在我的 Jetson Nano 上使用了大约 1.7GB 的 RAM。--cpu
使用该选项运行相同的脚本会给出:
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)
有没有办法减少这种开销?在 github issue 中提到了在没有所有 CUDA 内核的情况下编译 pytorch 以减少 RAM 开销,但我不确定我需要哪些编译选项以及哪些编译选项实际上会减少 RAM 开销。
有没有一种已知的方法可以减少 Pytorch 的 RAM 使用量?
解决方案
推荐阅读
- javascript - 将 addListener 用于稍后将创建的元素
- javascript - 如何以角度重置css
- signalr - AspNetCore / SignalR - 从 SignalR 集线器服务器返回 HttpResponseMessage 导致 InvalidDataException
- arrays - 仅当 counter_name 匹配时如何从 JSON 数据下方提取 counterid 并将 counterid 添加到 shell 脚本中的另一个文件
- ssh - 两个 Mercurial 回购。在远程机器上——一个可以通过 ssh 克隆;另一个,不是
- laravel - 如何在 Laravel Vue 中沿侧引导程序使用 zurb 基础?
- flutter - 如何在 SliverChildListDelegate 中使用 FirestoreAnimatedList?
- android - E/CRASH(7870):信号 11 (SIGSEGV),代码 1 (SEGV_MAPERR),故障地址 0000007c2f67f818
- amazon-web-services - AWS ASG 冷却期和扩展策略中的预热期有什么区别?
- javascript - 在 highcharts 散点图中添加自定义变量作为工具提示