使用 CUDA 时如何减少 PyTorch 使用的（非 GPU）RAM 开销？

Question

我正在尝试在我的 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 问题中，提到在没有所有 CUDA 内核的情况下编译 pytorch 以减少 RAM 开销，但我不确定我需要哪些编译选项以及哪些实际上会减少 RAM 开销。

有没有已知的方法可以减少 PyTorch 的 RAM 使用量？