多少人用 PyTorch“炼丹”时都会被这个 bug 困扰。
CUDA error: out of memory.[/pre]
一般情况下，你得找出当下占显存的没用的程序，然后 kill 掉。如果不行，还需手动调整 batch size 到合适的大小，有点麻烦。
现在，有人写了一个 PyTorch wrapper，
用一行代码就能“无痛”消除这个 bug。
有多厉害？
相关项目在 GitHub 才发布没几天就收获了 600 + 星。
一行代码解决内存溢出错误软件包名叫 koila，已经上传 PyPI，先安装一下：
pip install koila[/pre]现在，假如你面对这样一个 PyTorch 项目：构建一个神经网络来对 FashionMNIST 数据集中的图像进行分类。
先定义 input、label 和 model：
# A batch of MNIST imageinput = torch.randn(8, 28, 28) # A batch of labelslabel = 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[/pre]然后定义 loss 函数、计算输出和 losses。
loss_fn = CrossEntropyLoss() # Calculate losses out = nn(t)loss = loss_fn(out, label) # Backward passnn.zero_grad()loss.backward()[/pre]好了，如何使用 koila 来防止内存溢出？
超级简单！
只需在第一行代码，
也就是把输入用 lazy 张量 wrap 起来，并指定 bacth 维度，koila 就能自动帮你计算剩余的 GPU 内存并使用正确的 batch size 了。
在本例中，batch=0，则修改如下：
input = lazy(torch.randn(8, 28, 28), batch=0)[/pre]完事儿！就这样和 PyTorch“炼丹”时的 OOM 报错说拜拜。
灵感来自 TensorFlow 的静态 / 懒惰评估下面就来说说 koila 背后的工作原理。
“CUDA error: out of memory”这个报错通常发生在前向传递（forward pass）中，因为这时需要保存很多临时变量。
koila 的灵感来自 TensorFlow 的静态 / 懒惰评估（static / lazy evaluation）。
它通过构建图，并仅在必要时运行访问所有相关信息，来确定模型真正需要多少资源。
而只需计算临时变量的 shape 就能计算各变量的内存使用情况；而知道了在前向传递中使用了多少内存，koila 也就能自动选择最佳 batch size 了。
又是算 shape 又是算内存的，koila 听起来就很慢？
NO。
