Pytorch中torch.utils.checkpoint()

在PyTorch中，torch.utils.checkpoint 模塊提供了實(shí)現(xiàn)梯度檢查點(diǎn)（也稱(chēng)為checkpointing）的功能。這個(gè)技術(shù)主要用于訓(xùn)練時(shí)內(nèi)存優(yōu)化，它允許我們以計(jì)算時(shí)間為代價(jià)，減少訓(xùn)練深度網(wǎng)絡(luò)時(shí)的內(nèi)存占用。

原理

梯度檢查點(diǎn)技術(shù)的基本原理是，在前向傳播的過(guò)程中，并不保存所有的中間激活值。相反，它只保存一部分關(guān)鍵的激活值。在反向傳播時(shí)，根據(jù)保留的激活值重新計(jì)算丟棄的中間激活值。因此內(nèi)存的使用量會(huì)下降，但計(jì)算量會(huì)增加，因?yàn)樾枰匦掠?jì)算一些前向傳播的部分。

用法

torch.utils.checkpoint 中主要的函數(shù)是 checkpoint。checkpoint 函數(shù)可以用來(lái)封裝模型的一部分或者一個(gè)復(fù)雜的運(yùn)算，這部分會(huì)使用梯度檢查點(diǎn)。它的一般用法是：

import torch
from torch.utils.checkpoint import checkpoint
# 定義一個(gè)前向傳播函數(shù)
def custom_forward(*inputs):
    # 定義你的前向傳播邏輯
    # 例如: x, y = inputs; result = x + y
    ...
    return result
# 在訓(xùn)練的前向傳播過(guò)程中使用梯度檢查點(diǎn)
model_output = checkpoint(custom_forward, *model_inputs)

在每次調(diào)用 custom_forward 函數(shù)時(shí)，它都會(huì)返回正常的前向傳播結(jié)果。不過(guò)，checkpoint 函數(shù)會(huì)確保僅保留必須的激活值（即 custom_forward 的輸出）。其他激活值不會(huì)保存在內(nèi)存中，需要在反向傳播時(shí)重新計(jì)算。

下面是一個(gè)具體的示例，演示了如何在一個(gè)簡(jiǎn)單的模型中使用 checkpoint 函數(shù)：

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint
class SomeModel(nn.Module):
    def __init__(self):
        super(SomeModel, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 50, 5)
    def forward(self, x):
        # 使用checkpoint來(lái)減少第二層卷積的內(nèi)存使用量
        x = self.conv1(x)
        x = checkpoint(self.conv2, x)
        return x
model = SomeModel()
input = torch.randn(1, 1, 28, 28)
output = model(input)
loss = output.sum()
loss.backward()

在上面的例子中，conv2的前向計(jì)算是通過(guò) checkpoint 封裝的，這意味著在 conv1 的輸出和 conv2 的輸出之間的激活值不會(huì)被完全存儲(chǔ)。在反向傳播時(shí)，這些丟失的激活值會(huì)通過(guò)再次前向傳遞 conv2 來(lái)重新計(jì)算。
使用梯度檢查點(diǎn)技術(shù)可以在訓(xùn)練大型模型時(shí)減少顯存的占用，但由于在反向傳播時(shí)額外的重新計(jì)算，它會(huì)增加一些計(jì)算成本。

到此這篇關(guān)于Pytorch中torch.utils.checkpoint()及用法詳解的文章就介紹到這了,更多相關(guān)Pytorch torch.utils.checkpoint()內(nèi)容請(qǐng)搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！

最新評(píng)論