🐨 Koila

Koila 可以无痛解决 CUDA error: out of memory error 问题。
只需一行代码即可修复，然后忘掉它。

Koila

🚨 警告

主分支是项目的一个全新重构（目前因为我没有足够时间完成，所以大部分是空的）。要查看可用代码，请检出 v0.1.1 标签的概念验证版本（不完全支持所有操作，不适合生产环境）。要使用它，请下载 v0.1.1 版本。

🚀 功能

🙅 通过一行代码防止 CUDA error: out of memory error。
⚗️ 当批次大小过大时自动累积梯度。
🦥 懒加载执行 PyTorch 代码以节省计算资源。
✂️ 自动沿批次维度拆分为更适合 GPU 的数量（2 的幂）以加快执行速度。
🤏 最小化 API（只需包装所有输入即可）。

🤔 为什么选择 Koila？

是否曾遇到 RuntimeError: CUDA error: out of memory？
我们都喜欢 PyTorch 的速度、效率和透明性，但这也意味着它不会做额外的事情。比如，防止自 2017 年以来困扰许多用户的一个非常常见的错误（2017）。

这个库旨在通过对原生 PyTorch 进行轻量级包装来防止这种情况发生。当一个张量被包装时，库会自动计算剩余的 GPU 内存，并使用合适的批次大小，免去了每次使用模型时手动调整批次大小的麻烦。

此外，库会自动为 GPU 选择合适的批次大小。你知道使用更大的批次并不总是加快处理速度吗？这个库也会自动处理这一点。

因为 Koila 的代码其实就是 PyTorch 代码，它在底层运行 PyTorch，所以你可以一起使用而不必担心兼容性问题。

哦，而且这一切只需一行代码！😊

⬇️ 安装

Koila 已在 PyPI 上发布。要安装，请运行以下命令：

pip install koila

🏃 入门

使用非常简单。例如，你有如下 PyTorch 代码（来自 PyTorch 的教程）：

定义输入、标签和模型：

# 一批 MNIST 图片
input = torch.randn(8, 28, 28)

# 一批标签
label = torch.randn(0, 10, [8])

class NeuralNetwork(Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = Flatten()
        self.linear_relu_stack = Sequential(
            Linear(28 * 28, 512),
            ReLU(),
            Linear(512, 512),
            ReLU(),
            Linear(512, 10),
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

定义损失函数，计算输出和损失：

loss_fn = CrossEntropyLoss()

# 计算损失
out = nn(t)
loss = loss_fn(out, label)

# 反向传播
nn.zero_grad()
loss.backward()

好吧，如何将代码适配使用 Koila 的功能？

你只需添加这一行代码（v0.1.1 版本）：

# 包装输入张量和标签张量。
# 如果提供了批次参数，该张量的维度将被视为批次维度。
# 在此情况下，第一个维度（dim=0）用作批次维度。
(input, label) = lazy(input, label, batch=0)

搞定。你将不会再遇到内存不足的问题。

🏋️ 它的工作原理是什么？

CUDA error: out of memory 通常发生在前向传播时，因为临时变量需要保存在内存中。

Koila 是一个轻量级的 PyTorch 包装库。它受到 TensorFlow 的静态/惰性评估的启发。通过先构建图表，并仅在必要时运行模型，模型可以访问所有必要的信息来确定计算模型实际需要多少资源。

在内存使用方面，仅需要临时变量的形状即可计算模型中使用的这些变量的内存使用量。例如，+ 操作接收两个相同大小的张量，并输出一个大小等于输入大小的张量，而 log 操作接收一个张量，并输出另一个形状相同的张量。广播使得事情稍微复杂了一点，但总体思想是相同的。通过跟踪所有这些形状，人们可以轻松地判断前向传播中使用了多少内存，并相应地选择最佳的批次大小。