pytorch版本densenet代碼講解

DenseNet 模型代碼詳解

下面是 DenseNet 模型代碼的逐部分詳細解析：

1. 導入模塊

import re
from collections import OrderedDict
from functools import partial
from typing import Any, Optionalimport torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as cp
from torch import Tensorfrom ..transforms._presets import ImageClassification
from ..utils import _log_api_usage_once
from ._api import register_model, Weights, WeightsEnum
from ._meta import _IMAGENET_CATEGORIES
from ._utils import _ovewrite_named_param, handle_legacy_interface

re: 正則表達式模塊，用于處理權重名稱的轉換
OrderedDict: 有序字典，用于按順序構建網絡層
partial: 創建部分函數，用于預設圖像轉換參數
torch.nn: PyTorch 的神經網絡模塊
torch.utils.checkpoint: 內存優化技術，減少訓練時的內存占用
ImageClassification: 圖像分類的預處理轉換
register_model: 注冊模型的裝飾器
Weights/WeightsEnum: 預訓練權重相關類
_IMAGENET_CATEGORIES: ImageNet 數據集類別標簽
模型工具函數: 覆蓋參數、處理舊版接口等

2. DenseNet 基礎層 (_DenseLayer)

class _DenseLayer(nn.Module):def __init__(self, num_input_features: int, growth_rate: int, bn_size: int, drop_rate: float, memory_efficient: bool = False) -> None:super().__init__()# 第一個卷積塊 (1x1 卷積)self.norm1 = nn.BatchNorm2d(num_input_features)self.relu1 = nn.ReLU(inplace=True)self.conv1 = nn.Conv2d(num_input_features, bn_size * growth_rate, kernel_size=1, stride=1, bias=False)# 第二個卷積塊 (3x3 卷積)self.norm2 = nn.BatchNorm2d(bn_size * growth_rate)self.relu2 = nn.ReLU(inplace=True)self.conv2 = nn.Conv2d(bn_size * growth_rate, growth_rate, kernel_size=3, stride=1, padding=1, bias=False)self.drop_rate = float(drop_rate)self.memory_efficient = memory_efficient

Bottleneck 結構: 由兩個卷積層組成，減少計算量
1x1 卷積: 降維，輸出通道數為 bn_size * growth_rate
3x3 卷積: 主卷積層，輸出通道數為 growth_rate
memory_efficient: 是否使用梯度檢查點節省內存

前向傳播邏輯

    def bn_function(self, inputs: list[Tensor]) -> Tensor:# 拼接所有輸入特征concated_features = torch.cat(inputs, 1)# 通過第一個卷積塊bottleneck_output = self.conv1(self.relu1(self.norm1(concated_features)))return bottleneck_outputdef forward(self, input: Tensor) -> Tensor:if isinstance(input, Tensor):prev_features = [input]else:prev_features = input# 內存高效模式處理if self.memory_efficient and self.any_requires_grad(prev_features):bottleneck_output = self.call_checkpoint_bottleneck(prev_features)else:bottleneck_output = self.bn_function(prev_features)# 通過第二個卷積塊new_features = self.conv2(self.relu2(self.norm2(bottleneck_output)))# 應用Dropoutif self.drop_rate > 0:new_features = F.dropout(new_features, p=self.drop_rate, training=self.training)return new_features

特征拼接: 將前面所有層的輸出拼接在一起
梯度檢查點: 在內存高效模式下，使用檢查點減少內存占用
Dropout: 隨機丟棄部分神經元，防止過擬合

3. Dense 塊 (_DenseBlock)

class _DenseBlock(nn.ModuleDict):def __init__(self,num_layers: int,num_input_features: int,bn_size: int,growth_rate: int,drop_rate: float,memory_efficient: bool = False,) -> None:super().__init__()# 創建多個密集層for i in range(num_layers):layer = _DenseLayer(num_input_features + i * growth_rate,growth_rate=growth_rate,bn_size=bn_size,drop_rate=drop_rate,memory_efficient=memory_efficient,)self.add_module("denselayer%d" % (i + 1), layer)

模塊字典: 存儲多個密集層
輸入特征計算: 每增加一層，輸入特征增加 growth_rate 個通道

前向傳播

    def forward(self, init_features: Tensor) -> Tensor:features = [init_features]# 逐層處理并收集輸出for name, layer in self.items():new_features = layer(features)features.append(new_features)# 拼接所有層的輸出return torch.cat(features, 1)

特征累積: 每一層的輸出都添加到特征列表中
特征拼接: 將所有層的輸出沿通道維度拼接

4. 過渡層 (_Transition)

class _Transition(nn.Sequential):def __init__(self, num_input_features: int, num_output_features: int) -> None:super().__init__()# 壓縮特征維度self.norm = nn.BatchNorm2d(num_input_features)self.relu = nn.ReLU(inplace=True)self.conv = nn.Conv2d(num_input_features, num_output_features, kernel_size=1, stride=1, bias=False)# 空間下采樣self.pool = nn.AvgPool2d(kernel_size=2, stride=2)

特征壓縮: 1x1 卷積減少通道數（通常減半）
空間降維: 平均池化減小特征圖尺寸

5. DenseNet 主模型

class DenseNet(nn.Module):def __init__(self,growth_rate: int = 32,block_config: tuple[int, int, int, int] = (6, 12, 24, 16),num_init_features: int = 64,bn_size: int = 4,drop_rate: float = 0,num_classes: int = 1000,memory_efficient: bool = False,) -> None:super().__init__()_log_api_usage_once(self)  # 記錄API使用情況# 初始卷積層self.features = nn.Sequential(OrderedDict([("conv0", nn.Conv2d(3, num_init_features, kernel_size=7, stride=2, padding=3, bias=False)),("norm0", nn.BatchNorm2d(num_init_features)),("relu0", nn.ReLU(inplace=True)),("pool0", nn.MaxPool2d(kernel_size=3, stride=2, padding=1)),]))# 構建多個Dense塊和過渡層num_features = num_init_featuresfor i, num_layers in enumerate(block_config):# 添加Dense塊block = _DenseBlock(num_layers=num_layers,num_input_features=num_features,bn_size=bn_size,growth_rate=growth_rate,drop_rate=drop_rate,memory_efficient=memory_efficient,)self.features.add_module("denseblock%d" % (i + 1), block)num_features += num_layers * growth_rate# 添加過渡層（最后一個塊除外）if i != len(block_config) - 1:trans = _Transition(num_features, num_features // 2)self.features.add_module("transition%d" % (i + 1), trans)num_features = num_features // 2# 最終批歸一化self.features.add_module("norm5", nn.BatchNorm2d(num_features))# 分類器self.classifier = nn.Linear(num_features, num_classes)# 參數初始化for m in self.modules():if isinstance(m, nn.Conv2d):nn.init.kaiming_normal_(m.weight)elif isinstance(m, nn.BatchNorm2d):nn.init.constant_(m.weight, 1)nn.init.constant_(m.bias, 0)elif isinstance(m, nn.Linear):nn.init.constant_(m.bias, 0)

初始卷積層: 快速下采樣輸入圖像
塊配置: 控制每個Dense塊中的層數
通道管理: 通過過渡層壓縮通道數
Kaiming初始化: 卷積層的權重初始化
批歸一化初始化: 權重設為1，偏置設為0

前向傳播

    def forward(self, x: Tensor) -> Tensor:features = self.features(x)out = F.relu(features, inplace=True)out = F.adaptive_avg_pool2d(out, (1, 1))  # 全局平均池化out = torch.flatten(out, 1)  # 展平特征out = self.classifier(out)  # 分類return out

特征提取: 通過多個Dense塊和過渡層
全局平均池化: 將特征圖轉換為特征向量
全連接層: 輸出分類結果

6. 權重加載函數

def _load_state_dict(model: nn.Module, weights: WeightsEnum, progress: bool) -> None:# 匹配舊版權重名稱模式pattern = re.compile(r"^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$")state_dict = weights.get_state_dict(progress=progress, check_hash=True)# 轉換權重名稱for key in list(state_dict.keys()):res = pattern.match(key)if res:new_key = res.group(1) + res.group(2)state_dict[new_key] = state_dict[key]del state_dict[key]# 加載權重model.load_state_dict(state_dict)

權重名稱轉換: 適配舊版權重命名方式
哈希校驗: 確保下載的權重文件完整無誤

7. 模型工廠函數

def _densenet(growth_rate: int,block_config: tuple[int, int, int, int],num_init_features: int,weights: Optional[WeightsEnum],progress: bool,**kwargs: Any,
) -> DenseNet:# 根據權重調整輸出類別數if weights is not None:_ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"]))# 創建模型model = DenseNet(growth_rate, block_config, num_init_features, **kwargs)# 加載預訓練權重if weights is not None:_load_state_dict(model=model, weights=weights, progress=progress)return model

參數覆蓋: 根據預訓練權重調整輸出類別數
靈活配置: 支持不同DenseNet變體

8. 預訓練權重定義

_COMMON_META = {"min_size": (29, 29),  # 最小輸入尺寸"categories": _IMAGENET_CATEGORIES,  # ImageNet類別"recipe": "https://github.com/pytorch/vision/pull/116",  # 訓練方法
}class DenseNet121_Weights(WeightsEnum):IMAGENET1K_V1 = Weights(url="https://download.pytorch.org/models/densenet121-a639ec97.pth",transforms=partial(ImageClassification, crop_size=224),  # 圖像預處理meta={**_COMMON_META,"num_params": 7978856,  # 參數量"_metrics": {  # 性能指標"ImageNet-1K": {"acc@1": 74.434,  # top-1準確率"acc@5": 91.972,  # top-5準確率}},"_ops": 2.834,  # 計算量 (GFLOPs)"_file_size": 30.845,  # 文件大小 (MB)},)DEFAULT = IMAGENET1K_V1  # 默認權重

權重元數據: 包含模型性能和資源信息
預處理定義: 指定圖像分類任務的預處理流程
性能指標: 提供在ImageNet上的評估結果

9. 模型變體實現

@register_model()  # 注冊模型
@handle_legacy_interface(weights=("pretrained", DenseNet121_Weights.IMAGENET1K_V1))
def densenet121(*, weights: Optional[DenseNet121_Weights] = None, progress: bool = True, **kwargs: Any) -> DenseNet:weights = DenseNet121_Weights.verify(weights)  # 驗證權重return _densenet(32, (6, 12, 24, 16), 64, weights, progress, **kwargs)

DenseNet121: 增長率32，塊配置[6,12,24,16]，初始特征64
DenseNet169: 增長率32，塊配置[6,12,32,32]，初始特征64
DenseNet201: 增長率32，塊配置[6,12,48,32]，初始特征64
DenseNet161: 增長率48，塊配置[6,12,36,24]，初始特征96