終于把深度學(xué)習(xí)中的模型壓縮搞懂了！

作者：程序員小寒 2025-02-17 13:09:59

在深度學(xué)習(xí)中，模型壓縮是減少模型大小、降低計(jì)算復(fù)雜度，同時(shí)盡可能保持模型性能的一類技術(shù)。它在移動(dòng)端、嵌入式設(shè)備和邊緣計(jì)算等資源受限的環(huán)境中尤其重要。

今天給大家分享幾種常見的模型壓縮技術(shù)。

修剪

修剪是通過去除神經(jīng)網(wǎng)絡(luò)中某些不重要的連接或神經(jīng)元來(lái)減少模型的規(guī)模和計(jì)算需求。

修剪的目標(biāo)是去除那些對(duì)網(wǎng)絡(luò)性能影響較小的參數(shù)，從而達(dá)到減少模型復(fù)雜度的效果。

常見修剪策略

權(quán)重修剪
通過移除那些對(duì)網(wǎng)絡(luò)輸出貢獻(xiàn)較小的權(quán)重來(lái)減少模型的大小。
這些權(quán)重可以通過設(shè)定一個(gè)閾值來(lái)判定：低于某個(gè)閾值的權(quán)重會(huì)被剪掉。
神經(jīng)元修剪
修剪掉整個(gè)神經(jīng)元或通道，這樣的修剪方法可以進(jìn)一步減少計(jì)算量，尤其是對(duì)于卷積神經(jīng)網(wǎng)絡(luò)（CNN）來(lái)說(shuō)，移除不重要的特征圖通道會(huì)顯著降低計(jì)算復(fù)雜度。

修剪的步驟通常是：

訓(xùn)練原始模型
計(jì)算每個(gè)權(quán)重的重要性或每個(gè)神經(jīng)元的激活度
去除不重要的權(quán)重或神經(jīng)元
重新訓(xùn)練，以恢復(fù)性能損失

優(yōu)缺點(diǎn)

優(yōu)點(diǎn)：減小模型尺寸，降低計(jì)算負(fù)擔(dān)，提升推理速度，尤其適合硬件加速。
缺點(diǎn)：修剪過度可能導(dǎo)致模型性能下降。需要精心設(shè)計(jì)修剪方案，以在壓縮和性能之間找到平衡。

import torch
import torch.nn as nn
import torch.nn.utils.prune as prune

# 定義一個(gè)簡(jiǎn)單的神經(jīng)網(wǎng)絡(luò)
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(10, 10)
        self.fc2 = nn.Linear(10, 2)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 創(chuàng)建網(wǎng)絡(luò)和輸入數(shù)據(jù)
model = SimpleNet()
input_data = torch.randn(1, 10)

# 修剪fc1層的20%的權(quán)重
prune.random_unstructured(model.fc1, name="weight", amount=0.2)

# 打印fc1層的權(quán)重，觀察被修剪掉的權(quán)重
print(model.fc1.weight)

量化

量化是將浮點(diǎn)數(shù)表示的參數(shù)（如權(quán)重和激活）轉(zhuǎn)換為低精度數(shù)值表示（如整數(shù)）。

量化通常將模型從 32 位浮動(dòng)點(diǎn)數(shù)轉(zhuǎn)換為更低精度的數(shù)據(jù)類型，如16位、8位或更低，這樣可以減少存儲(chǔ)需求和加速推理過程。

權(quán)重量化
將模型中的浮點(diǎn)數(shù)權(quán)重轉(zhuǎn)換為低精度整數(shù)。例如，將32位浮點(diǎn)數(shù)權(quán)重映射到8位整數(shù)，這樣就能大幅減少模型的存儲(chǔ)需求。
激活量化
對(duì)于激活值（神經(jīng)網(wǎng)絡(luò)各層的輸出），也可以應(yīng)用類似的量化策略。

常見的量化類型

后訓(xùn)練量化：在模型訓(xùn)練完成后進(jìn)行量化，適用于已經(jīng)訓(xùn)練好的模型。
量化感知訓(xùn)練：在訓(xùn)練過程中加入量化過程，從而使得模型能夠適應(yīng)低精度的計(jì)算。

量化不僅減小了模型大小，還可能加速模型的推理過程，尤其是在支持低精度計(jì)算的硬件上（如TPU、GPU等）。

優(yōu)缺點(diǎn)

優(yōu)點(diǎn)：大幅減少模型的存儲(chǔ)需求，加速推理過程。尤其在嵌入式設(shè)備和移動(dòng)端設(shè)備上具有顯著的優(yōu)勢(shì)。
缺點(diǎn)：量化可能導(dǎo)致一定的精度損失

import torch
import torch.nn as nn
import torch.quantization

# 定義一個(gè)簡(jiǎn)單的模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(28 * 28, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 初始化模型
model = SimpleModel()

# 轉(zhuǎn)換模型為量化版本
model.eval()  # 切換到評(píng)估模式
quantized_model = torch.quantization.quantize_dynamic(
    model, {nn.Linear}, dtype=torch.qint8
)

# 查看量化后的模型
print(quantized_model)

蒸餾

蒸餾是一種將大型、復(fù)雜模型的知識(shí)遷移到較小模型中的技術(shù)。

通常，蒸餾的過程是訓(xùn)練一個(gè)小模型（學(xué)生模型）以模仿一個(gè)較大的、預(yù)先訓(xùn)練好的模型（教師模型）的行為。

小模型通過學(xué)習(xí)教師模型的預(yù)測(cè)概率分布來(lái)獲取知識(shí)，而不僅僅是傳統(tǒng)的標(biāo)簽信息。

蒸餾的主要思想是：

教師模型輸出的類別概率包含了更多的“軟信息”，這些信息能夠幫助學(xué)生模型更好地學(xué)習(xí)一些復(fù)雜的模式。
學(xué)生模型通過與教師模型輸出的“軟標(biāo)簽”進(jìn)行學(xué)習(xí)，能夠在不完全依賴硬標(biāo)簽的情況下獲取更多的信息，進(jìn)而提高其性能。

蒸餾的步驟通常是：

訓(xùn)練教師模型
首先，訓(xùn)練一個(gè)大型且高性能的教師模型，這通常是一個(gè)深度神經(jīng)網(wǎng)絡(luò)。
訓(xùn)練學(xué)生模型
然后，訓(xùn)練一個(gè)較小的學(xué)生模型，目標(biāo)是通過最小化學(xué)生模型與教師模型在相同輸入上的輸出差異來(lái)進(jìn)行訓(xùn)練。學(xué)生模型不僅學(xué)習(xí)硬標(biāo)簽（真實(shí)標(biāo)簽），還學(xué)習(xí)教師模型的“軟標(biāo)簽”。

通過蒸餾，學(xué)生模型可以獲得教師模型中蘊(yùn)含的豐富知識(shí)，尤其是在教師模型能夠捕獲的復(fù)雜特征和模式方面，從而在保持較小規(guī)模的同時(shí)接近或達(dá)到教師模型的性能。

優(yōu)缺點(diǎn)

優(yōu)點(diǎn)：蒸餾可以顯著提高小型模型的性能，使其在壓縮后依然保持接近教師模型的精度，尤其在大型模型壓縮時(shí)表現(xiàn)出色。
缺點(diǎn)：蒸餾的一個(gè)挑戰(zhàn)是教師模型的選擇和訓(xùn)練需要耗費(fèi)大量的計(jì)算資源和時(shí)間。此外，蒸餾的效果可能會(huì)受到學(xué)生模型的限制，對(duì)于某些任務(wù)，學(xué)生模型的性能可能不容易達(dá)到教師模型的水平。

import torch
import torch.nn as nn
import torch.optim as optim

# 創(chuàng)建一個(gè)簡(jiǎn)單的教師模型和學(xué)生模型
class TeacherNet(nn.Module):
    def __init__(self):
        super(TeacherNet, self).__init__()
        self.fc = nn.Linear(10, 10)

    def forward(self, x):
        return self.fc(x)

class StudentNet(nn.Module):
    def __init__(self):
        super(StudentNet, self).__init__()
        self.fc = nn.Linear(10, 10)

    def forward(self, x):
        return self.fc(x)

# 創(chuàng)建模型實(shí)例
teacher = TeacherNet()
student = StudentNet()

# 使用教師模型生成“軟標(biāo)簽”
def distillation_loss(student_outputs, teacher_outputs, temperature=2.0):
    # 使用溫度縮放進(jìn)行蒸餾損失計(jì)算
    loss = nn.KLDivLoss()(nn.functional.log_softmax(student_outputs / temperature, dim=1),
                          nn.functional.softmax(teacher_outputs / temperature, dim=1)) * (temperature ** 2)
    return loss

# 簡(jiǎn)單的訓(xùn)練循環(huán)
optimizer = optim.SGD(student.parameters(), lr=0.1)

# 模擬訓(xùn)練過程
for epoch in range(100):
    # 輸入數(shù)據(jù)
    inputs = torch.randn(32, 10)  # 假設(shè)批次大小是32，輸入維度是10
    teacher_outputs = teacher(inputs)
    
    # 學(xué)生模型的輸出
    student_outputs = student(inputs)
    
    # 計(jì)算損失
    loss = distillation_loss(student_outputs, teacher_outputs)
    
    # 反向傳播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if epoch % 10 == 0:
        print(f"Epoch [{epoch}/100], Loss: {loss.item()}")

責(zé)任編輯：華軒來(lái)源：程序員學(xué)長(zhǎng)

深度學(xué)習(xí)模型壓縮量化