簡介面向深度學(xué)習(xí)初學(xué)者和圖像分類任務(wù)開發(fā)者該壓縮包提供了一份基于ResNet實現(xiàn)CIFAR-10分類的完整算法工程。代碼分為模型定義、數(shù)據(jù)加載、訓(xùn)練與推理四個Python腳本并附有說明文檔共五個文件體積僅三KB結(jié)構(gòu)精簡、無額外依賴便于快速閱讀和二次修改。當(dāng)前已有七百余人瀏覽學(xué)習(xí)。實現(xiàn)中重點涵蓋了殘差塊的跳躍連接、批量歸一化、全局平均池化、交叉熵?fù)p失以及帶動量的隨機梯度下降優(yōu)化器并引入了隨機裁剪、水平翻轉(zhuǎn)等數(shù)據(jù)增強手段。讀者可以對照訓(xùn)練腳本直觀理解ResNet如何通過恒等映射緩解深層網(wǎng)絡(luò)的梯度消失問題并掌握針對三十二乘三十二像素小尺寸圖像的分類管線設(shè)計方法。適合作為高校深度學(xué)習(xí)課程作業(yè)、競賽入門或論文復(fù)現(xiàn)的參考代碼。1. ResNet CIFAR10一個小數(shù)據(jù)集上的經(jīng)典分類實戰(zhàn)做深度學(xué)習(xí)圖像分類的人遲早都會跟 CIFAR10 打交道。這個數(shù)據(jù)集只有 10 個類別、6 萬張 32×32 的小圖但想在它上面把準(zhǔn)確率推到 90% 以上并不是隨便疊幾層卷積就能做到的事。這份資源不是簡單的“調(diào)庫跑通”而是一個從數(shù)據(jù)加載、模型搭建到訓(xùn)練調(diào)參都完整落地的 ResNet 實現(xiàn)——你拿到的是一整個能跑出結(jié)果的算法工程而不是一段粘貼即用的碎片代碼。它適合兩類人一類是剛學(xué)完卷積神經(jīng)網(wǎng)絡(luò)、想看看殘差連接到底怎么解決網(wǎng)絡(luò)退化問題的新手另一類是想在自己的分類任務(wù)里快速遷移 ResNet 結(jié)構(gòu)的開發(fā)者。下文會按“它怎么設(shè)計 → 核心代碼長什么樣 → 怎么跑起來 → 踩了哪些坑 → 怎么改成自己的模型”的順序把這份源碼逐層拆開。2. 項目架構(gòu)與數(shù)據(jù)流從 32×32 輸入到 10 類輸出2.1 目錄結(jié)構(gòu)一覽模型、訓(xùn)練、數(shù)據(jù)三塊分離拿到壓縮包解壓之后先別急著跑花兩分鐘把目錄結(jié)構(gòu)過一遍。這個項目的組織方式是典型的“單腳本訓(xùn)練 模塊化模型定義”比把幾百行代碼塞進(jìn)一個 main.py 要清晰得多。resnet-cifar10/ ├── model/ │ ├── resnet.py # ResNet 網(wǎng)絡(luò)結(jié)構(gòu)定義 │ └── __init__.py ├── data/ │ ├── dataset.py # CIFAR10 數(shù)據(jù)加載與增強 │ └── __init__.py ├── train.py # 訓(xùn)練主腳本 ├── test.py # 測試與評估腳本 ├── config.py # 全局超參數(shù)配置 └── checkpoints/ # 模型權(quán)重保存目錄運行時生成我一般拿到別人的項目會先看 config.py因為超參數(shù)往往是整個工程最“玄學(xué)”的部分。這個項目的 config.py 里定義了初始學(xué)習(xí)率、batch size、訓(xùn)練輪數(shù)、權(quán)重衰減系數(shù)還有數(shù)據(jù)增強開關(guān)。數(shù)據(jù)部分獨立成模塊意味著你后續(xù)想換成自己的數(shù)據(jù)集只需要改動 dataset.py 里的數(shù)據(jù)加載邏輯模型和訓(xùn)練腳本都可以不動。2.2 數(shù)據(jù)加載與增強CIFAR10 的預(yù)處理細(xì)節(jié)CIFAR10 的圖片分辨率只有 32×32這意味著不能直接套用 ImageNet 那套 224×224 的預(yù)處理流程。dataset.py 里做了幾個關(guān)鍵操作先將 PIL 圖像轉(zhuǎn)為 Tensor再做標(biāo)準(zhǔn)化三個通道的均值是 (0.4914, 0.4822, 0.4465)標(biāo)準(zhǔn)差是 (0.2023, 0.1994, 0.2010)。這一組數(shù)值是 CIFAR10 數(shù)據(jù)集的全局統(tǒng)計量是官方計算好的直接拿來用就行不要自己重新算否則會引入不必要的偏差。# data/dataset.py 核心代碼 transform_train transforms.Compose([ transforms.RandomCrop(32, padding4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize( mean[0.4914, 0.4822, 0.4465], std[0.2023, 0.1994, 0.2010] ) ]) transform_test transforms.Compose([ transforms.ToTensor(), transforms.Normalize( mean[0.4914, 0.4822, 0.4465], std[0.2023, 0.1994, 0.2010] ) ])這段代碼里訓(xùn)練集和測試集的預(yù)處理差異是刻意設(shè)計的訓(xùn)練集做了 RandomCrop 和 RandomHorizontalFlip屬于數(shù)據(jù)增強目的是讓模型在有限的 5 萬張訓(xùn)練圖上學(xué)到更魯棒的特征測試集只做標(biāo)準(zhǔn)化不做任何隨機變換保證評估結(jié)果的可復(fù)現(xiàn)性。RandomCrop 的 padding4 表示在圖片四周填充 4 個像素后再隨機裁剪回 32×32這是 CIFAR10 上最經(jīng)典也最有效的增強手段幾乎沒有之一。2.3 訓(xùn)練主循環(huán)學(xué)習(xí)率調(diào)度與梯度更新train.py 里的訓(xùn)練循環(huán)寫得比較規(guī)矩每個 epoch 遍歷訓(xùn)練集計算損失、反向傳播、更新權(quán)重然后每個 epoch 結(jié)束在驗證集上評估一次。值得留意的是學(xué)習(xí)率調(diào)度策略這個項目用了余弦退火CosineAnnealingLR初始學(xué)習(xí)率 0.1隨著訓(xùn)練輪數(shù)逐步衰減到接近 0。相比固定學(xué)習(xí)率余弦退火在 CIFAR10 上通常能多漲 12 個點的準(zhǔn)確率屬于性價比很高的改動。# train.py 訓(xùn)練核心邏輯 for epoch in range(config.epochs): model.train() for batch_idx, (inputs, targets) in enumerate(trainloader): inputs, targets inputs.to(device), targets.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, targets) loss.backward() optimizer.step() scheduler.step() # 每個 epoch 結(jié)束后在測試集上評估 accuracy evaluate(model, testloader, device) print(fEpoch {epoch1}/{config.epochs}, Acc: {accuracy:.2f}%)這里的 device 變量通過config.device控制支持 GPU 和 CPU 兩種模式。如果機器上沒有 CUDA代碼會自動回退到 CPU 訓(xùn)練但 200 個 epoch 跑下來會比較慢建議還是在有 GPU 的環(huán)境下運行。criterion用的是 CrossEntropyLoss這是多分類任務(wù)的標(biāo)準(zhǔn)損失函數(shù)它內(nèi)部已經(jīng)包含了 Softmax 操作所以模型最后一層不需要額外加 Softmax。3. 核心算法拆解ResNet 的殘差連接與瓶頸塊實現(xiàn)3.1 BasicBlock兩層卷積的殘差學(xué)習(xí)單元這個項目采用的是 ResNet 的 CIFAR 版本變體網(wǎng)絡(luò)結(jié)構(gòu)比 ImageNet 版本要輕量得多。核心組件是 BasicBlock由兩個 3×3 卷積組成中間夾一個 BatchNorm 和 ReLU最后有一條恒等映射shortcut把輸入直接加到輸出上。殘差連接的意義在于當(dāng)網(wǎng)絡(luò)層數(shù)加深時恒等映射提供了一個“高速公路”讓梯度可以無損回傳從而緩解梯度消失問題。# model/resnet.py BasicBlock 實現(xiàn) class BasicBlock(nn.Module): def __init__(self, in_channels, out_channels, stride1): super(BasicBlock, self).__init__() self.conv1 nn.Conv2d(in_channels, out_channels, kernel_size3, stridestride, padding1, biasFalse) self.bn1 nn.BatchNorm2d(out_channels) self.conv2 nn.Conv2d(out_channels, out_channels, kernel_size3, stride1, padding1, biasFalse) self.bn2 nn.BatchNorm2d(out_channels) # 輸入輸出通道不一致時用 1x1 卷積對齊維度 self.shortcut nn.Sequential() if stride ! 1 or in_channels ! out_channels: self.shortcut nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size1, stridestride, biasFalse), nn.BatchNorm2d(out_channels) ) def forward(self, x): residual x out torch.relu(self.bn1(self.conv1(x))) out self.bn2(self.conv2(out)) out self.shortcut(residual) out torch.relu(out) return out這里的biasFalse是 BatchNorm 的搭配慣例——卷積層后面緊跟 BatchNorm 時偏置項會被 BN 層歸一化掉留著只會增加參數(shù)冗余。shortcut 分支的 1×1 卷積用于兩種場景一是 stride2 時特征圖尺寸減半需要 1×1 卷積同時降采樣并對齊通道數(shù)二是通道數(shù)翻倍時需要 1×1 卷積把恒等映射的通道數(shù)擴(kuò)到和主分支一致否則兩個張量無法相加。3.2 網(wǎng)絡(luò)組裝堆疊 BasicBlock 構(gòu)建深度模型整個 ResNet 的組裝邏輯在于各階段的下采樣。CIFAR 版本的 ResNet 輸入層通常先用一個 3×3 卷積而不是 ImageNet 的 7×7MaxPool因為 32×32 的輸入分辨率太小大卷積核會過早縮小特征圖。項目里默認(rèn)的是 ResNet20結(jié)構(gòu)是一層初始卷積 三層 BasicBlock 組每組分別堆疊 3、3、3 個塊通道數(shù)從 16 開始逐階段翻倍為 16→32→64。# model/resnet.py ResNet 主體結(jié)構(gòu) class ResNet(nn.Module): def __init__(self, block, num_blocks, num_classes10): super(ResNet, self).__init__() self.in_channels 16 self.conv1 nn.Conv2d(3, 16, kernel_size3, stride1, padding1, biasFalse) self.bn1 nn.BatchNorm2d(16) self.layer1 self._make_layer(block, 16, num_blocks[0], stride1) self.layer2 self._make_layer(block, 32, num_blocks[1], stride2) self.layer3 self._make_layer(block, 64, num_blocks[2], stride2) self.avg_pool nn.AdaptiveAvgPool2d((1, 1)) self.fc nn.Linear(64, num_classes) def _make_layer(self, block, out_channels, num_blocks, stride): strides [stride] [1] * (num_blocks - 1) layers [] for s in strides: layers.append(block(self.in_channels, out_channels, strides)) self.in_channels out_channels return nn.Sequential(*layers)_make_layer這個函數(shù)有一個容易看懵的細(xì)節(jié)每組第一個 block 的 stride 負(fù)責(zé)下采樣后面的 block 保持 stride1。比如 layer2 傳入 stride2那么這一層第一個 BasicBlock 的卷積 stride2特征圖從 16×16 降到 8×8同時 shortcut 也要做一次 stride2 的 1×1 卷積來匹配尺寸。后面的兩個 block 維持 8×8不斷提取更深層的語義特征。3.3 參數(shù)量對比為什么 CIFAR 版本要輕如果你拿這個網(wǎng)絡(luò)和標(biāo)準(zhǔn)的 ResNet50 對比會發(fā)現(xiàn)參數(shù)量差了一個數(shù)量級。ResNet50 的參數(shù)大約 2500 萬而這里的 ResNet20 只有約 27 萬參數(shù)。原因是 CIFAR10 的圖只有 32×32信息量遠(yuǎn)小于 ImageNet 的 224×224把網(wǎng)絡(luò)做得過深反而會過擬合。實際經(jīng)驗是ResNet20 在 CIFAR10 上就能達(dá)到 92% 左右的準(zhǔn)確率ResNet32 能到 93%再往上加到 ResNet56 收益就開始邊際遞減了訓(xùn)練時間卻成倍增長。網(wǎng)絡(luò)尾部用了AdaptiveAvgPool2d((1, 1))而不是傳統(tǒng)的 FlattenMaxPool這樣不管輸入分辨率怎么變?nèi)B接層的輸入維度始終是 64。這個設(shè)計在你遷移到其他尺寸的數(shù)據(jù)集時非常有用——比如你想改成 48×48 的輸入網(wǎng)絡(luò)結(jié)構(gòu)不用動只需要改數(shù)據(jù)加載部分。4. 運行與調(diào)參指南一次完整的訓(xùn)練流程復(fù)現(xiàn)4.1 環(huán)境準(zhǔn)備與依賴安裝在跑這個項目之前先確認(rèn)環(huán)境里有 Python 3.8 以上的解釋器以及 PyTorch 和 torchvision。我推薦用 conda 創(chuàng)建獨立環(huán)境避免把系統(tǒng) Python 搞亂。安裝命令如下# 創(chuàng)建虛擬環(huán)境并安裝依賴 conda create -n resnet-cifar10 python3.9 conda activate resnet-cifar10 pip install torch torchvision # 如果需要 GPU 版本按官網(wǎng)命令安裝 pip install numpy tqdm matplotlib如果你的機器有 NVIDIA GPU建議安裝 CUDA 版的 PyTorch訓(xùn)練速度會快 2030 倍。沒有 GPU 也沒關(guān)系這個項目在網(wǎng)絡(luò)和 batch size 上都做了 CPU 友好的設(shè)定只是 200 個 epoch 可能需要跑幾個小時耐心等就好。安裝完成之后在項目根目錄執(zhí)行python train.py腳本會自動下載 CIFAR10 數(shù)據(jù)集到本地首次運行需要聯(lián)網(wǎng)約 170MB。4.2 超參數(shù)含義與調(diào)整建議config.py 里的超參數(shù)不是隨手填的每個都對應(yīng)一個訓(xùn)練效果的關(guān)鍵變量。下面這個表列了最核心的幾個參數(shù)以及我的調(diào)參經(jīng)驗參數(shù)名默認(rèn)值含義調(diào)整建議epochs200訓(xùn)練輪數(shù)少于 100 輪模型欠擬合多于 300 輪收益極小batch_size128單次前向傳播的樣本數(shù)顯存不夠時降到 64同時可適當(dāng)降低學(xué)習(xí)率learning_rate0.1初始學(xué)習(xí)率batch size 減半時學(xué)習(xí)率也應(yīng)減半weight_decay5e-4L2 正則化系數(shù)過擬合加重時調(diào)大到 1e-3momentum0.9SGD 動量一般不動0.9 是默認(rèn)最優(yōu)值這里有一個新手常犯的錯誤直接套用 ImageNet 分類的初始學(xué)習(xí)率 0.01。在 CIFAR10 ResNet 的組合下0.01 的初始學(xué)習(xí)率會導(dǎo)致收斂極其緩慢因為 batch size 小、梯度噪聲大。項目默認(rèn)的 0.1 配合帶動量的 SGD前幾個 epoch 就能看到損失顯著下降。如果你想用 Adam 替代 SGD建議把學(xué)習(xí)率降到 1e-3因為 Adam 對學(xué)習(xí)率的敏感度更高。4.3 訓(xùn)練輸出解讀與斷點恢復(fù)訓(xùn)練過程中每個 epoch 都會打印準(zhǔn)確率和損失值。正常的訓(xùn)練曲線應(yīng)該是這樣的前 10 個 epoch 準(zhǔn)確率快速從 10%隨機猜測水平爬到 70% 左右之后增速放緩100 個 epoch 左右達(dá)到 90% 以上隨后進(jìn)入平臺期靠余弦退火把最后的幾個點“磨”出來。Epoch 1/200, Acc: 32.45%, Loss: 2.310 Epoch 10/200, Acc: 69.87%, Loss: 0.982 Epoch 50/200, Acc: 87.23%, Loss: 0.354 Epoch 100/200, Acc: 90.18%, Loss: 0.271 Epoch 200/200, Acc: 92.37%, Loss: 0.189如果你看到 Loss 一開始就不降或者波動劇烈多半是學(xué)習(xí)率問題。此外項目里還寫了斷點保存邏輯——每 20 個 epoch 保存一次 checkpoint存放在 checkpoints/ 目錄下。中途中斷后想繼續(xù)訓(xùn)練可以用python train.py --resume加載最近的權(quán)重不用從頭開始。5. 避坑指南訓(xùn)練 CIFAR10 時的五個常見問題5.1 準(zhǔn)確率卡在 85% 上不去現(xiàn)象訓(xùn)練 Loss 正常下降但測試集準(zhǔn)確率始終在 85% 左右徘徊超過 20 個 epoch 都沒有突破。原因這一般是數(shù)據(jù)增強不足導(dǎo)致的過擬合。CIFAR10 的訓(xùn)練集只有 5 萬張圖如果只做 ToTensor 和 Normalize模型的泛化能力會比較差。解決確認(rèn) train.py 里確實用上了 RandomCrop 和 RandomHorizontalFlip。如果已經(jīng)用了還卡住可以再加一個 Cutout隨機遮擋一小塊區(qū)域或者 MixUp 增強這兩個技巧在 CIFAR10 上往往能再漲 1~2 個點。5.2 訓(xùn)練 Loss 變成 NaN現(xiàn)象Loss 在某個 epoch 突然變成 NaN之后所有參數(shù)都跟著變成 NaN。原因?qū)W習(xí)率太大導(dǎo)致梯度爆炸。0.1 的初始學(xué)習(xí)率配合 SGD 在大多數(shù)情況沒事但如果你把 batch size 調(diào)小了比如 32而學(xué)習(xí)率沒變梯度更新的方差會變大容易把權(quán)重推到數(shù)值溢出的區(qū)域。解決檢查最近改了什么參數(shù)優(yōu)先把學(xué)習(xí)率降到 0.05 或者 0.01 重試。另一個技巧是給模型加梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm5)。5.3 從零訓(xùn)練比官方準(zhǔn)確率低了 3 個百分點現(xiàn)象訓(xùn)練完成后準(zhǔn)確率在 89% 左右但論文或者 README 里寫的 baseline 是 92%。原因這種差距絕大多數(shù)來自兩個地方一是訓(xùn)練輪數(shù)不夠論文里的 200 個 epoch 是有必要的二是學(xué)習(xí)率調(diào)度沒有生效。很多簡化項目的代碼里學(xué)的是 step decay每隔 60 個 epoch 衰減一次而這個項目用的是 CosineAnnealingLR兩者在同一網(wǎng)絡(luò)上的差距能到 1~2 個點。解決確認(rèn)scheduler.step()是在每個 epoch 結(jié)尾調(diào)用而不是在每次 batch 迭代里調(diào)用。用余弦退火時學(xué)習(xí)率是平滑下降的如果寫在錯誤的位置調(diào)度器會在一個 epoch 內(nèi)就把學(xué)習(xí)率耗盡模型后半程基本白學(xué)。5.4 測試集和訓(xùn)練集預(yù)處理不一致現(xiàn)象模型訓(xùn)練完測試準(zhǔn)確率很高95%但實際部署到業(yè)務(wù)場景里表現(xiàn)很差或者跑 test.py 時結(jié)果異常高。原因測試階段混入了數(shù)據(jù)增強。有的同學(xué)圖省事在 transform_test 里也加上了 RandomCrop 或 RandomHorizontalFlip導(dǎo)致評估結(jié)果虛高。解決test.py 里的 transform 必須和 transform_test 一致只包含 ToTensor 和 Normalize。另外驗證的時候用model.eval()關(guān)閉 dropout 和 BatchNorm 的更新否則 BN 層的統(tǒng)計量會亂掉。這是新手最容易忽略的細(xì)節(jié)。5.5 GPU 顯存不足現(xiàn)象batch_size128 時 CUDA out of memory。原因如果你改過網(wǎng)絡(luò)結(jié)構(gòu)比如把 BasicBlock 的通道數(shù)從 16 翻到 32參數(shù)量和中間特征圖的顯存占用會成倍增長。解決先把 batch_size 降到 64 或 32 試試。如果還想保持較大 batch可以用torch.cuda.amp混合精度訓(xùn)練顯存占用能降低約 40%。在這個項目里把model.half()和inputs.half()加上大部分 4GB 顯存的卡都能跑起來。6. 二次開發(fā)把 CIFAR10 分類器改成你的私有數(shù)據(jù)分類工具6.1 替換數(shù)據(jù)集換數(shù)據(jù)不換模型做二次開發(fā)最常見的需求是把這個分類工程遷移到自己的數(shù)據(jù)集上。CIFAR10 是 10 類的 32×32 小圖而你手頭的數(shù)據(jù)可能是貓狗分類2 類、工業(yè)缺陷檢測5 類或者某種特定場景的圖片。改動其實很小主要集中在兩個位置dataset.py 的數(shù)據(jù)加載函數(shù)和 model 最后一層的輸出維度。# 修改數(shù)據(jù)集加載為自定義圖片文件夾結(jié)構(gòu) # 目錄結(jié)構(gòu)data/train/class1/*.jpg, data/train/class2/*.jpg from torchvision import datasets train_dataset datasets.ImageFolder( rootdata/train, transformtransform_train ) test_dataset datasets.ImageFolder( rootdata/test, transformtransform_test )改完之后要注意兩個連帶問題一是num_classes要改成你自己的類別數(shù)比如貓狗分類就是 2只需要把ResNet(block, num_blocks, num_classes10)里的 10 改成 2二是全連接層的輸入尺寸不需要改因為AdaptiveAvgPool2d已經(jīng)把特征壓縮成固定 64 維了。如果你的圖片不是 32×32比如是 128×128就在 transform 里加一個transforms.Resize((32, 32))模型完全不用動。6.2 保存并復(fù)用模型從訓(xùn)練到推理訓(xùn)練好的模型不能只活在訓(xùn)練腳本里最后總歸要導(dǎo)出成可部署的形式。常見的做法是直接保存整個模型狀態(tài)字典然后在推理腳本里重新構(gòu)建網(wǎng)絡(luò)結(jié)構(gòu)再加載權(quán)重。# 訓(xùn)練完成后保存 torch.save(model.state_dict(), resnet20_cifar10.pth) # 推理時加載 model ResNet(BasicBlock, [3, 3, 3], num_classes10) model.load_state_dict(torch.load(resnet20_cifar10.pth)) model.eval() # 單張圖片預(yù)測 def predict(image_path): img Image.open(image_path).convert(RGB) img transform_test(img).unsqueeze(0) with torch.no_grad(): output model(img) pred output.argmax(dim1).item() return class_names[pred]這里model.eval()再強調(diào)一遍不是可選項。加載權(quán)重后如果不切換到 eval 模式BatchNorm 層的均值和方差還是會用訓(xùn)練時積累的統(tǒng)計量但 dropout 層的行為會保持訓(xùn)練模式導(dǎo)致推理結(jié)果不穩(wěn)定。上述代碼里no_grad()是另一個關(guān)鍵點它告訴 PyTorch 不要構(gòu)建計算圖節(jié)省內(nèi)存且加速推理。6.3 換個更大的網(wǎng)絡(luò)從 ResNet20 到 ResNet56當(dāng)你覺得當(dāng)前準(zhǔn)確率還不夠卷的時候增大模型容量是最直接的思路。項目里已經(jīng)預(yù)留了這個擴(kuò)展點——num_blocks參數(shù)就是用來控制深度的。把實例化代碼從ResNet(BasicBlock, [3, 3, 3])改成ResNet(BasicBlock, [5, 5, 5])就是 ResNet32改成[9, 9, 9]就是 ResNet56。# 在 config.py 中增加網(wǎng)絡(luò)深度配置 model ResNet(BasicBlock, [5, 5, 5], num_classes10) # ResNet32改網(wǎng)絡(luò)深度時其余超參數(shù)不用變但訓(xùn)練時間會顯著增加——ResNet32 大約是 ResNet20 的 1.6 倍。另外有個經(jīng)驗值網(wǎng)絡(luò)越深weight_decay 的作用越明顯。如果 ResNet56 的效果反而不如 ResNet32優(yōu)先檢查是不是過擬合了把 weight_decay 從 5e-4 調(diào)到 1e-3同時增加一點數(shù)據(jù)增強的強度。6.4 部署到 TorchScript 獲取加速最后一個技巧是把訓(xùn)練好的模型轉(zhuǎn)成 TorchScript這樣可以在沒有 Python 環(huán)境的 C 端部署或者通過 LibTorch 做推理加速。轉(zhuǎn)換過程很簡單但建議在 eval 模式下完成避免把 BN 層的訓(xùn)練行為固化進(jìn)腳本。# 導(dǎo)出 TorchScript 模型 script_model torch.jit.script(model) script_model.save(resnet20_torchscript.pt)我一般會把這一步放在模型最終確定之后。有個血淚教訓(xùn)是在 model.train() 模式下導(dǎo)出模型之后部署時輸出結(jié)果總是和驗證集差一點排查了很久才發(fā)現(xiàn)是 BN 層的 running_mean 和 running_var 被凍結(jié)成了非最優(yōu)狀態(tài)。從那以后我每次導(dǎo)出模型都強制先走一遍 eval 模式再驗證一整個測試集的準(zhǔn)確率沒掉點才放心。這份資源雖然針對的是 CIFAR10但它的代碼結(jié)構(gòu)和模塊劃分方式足夠撐起你后面大半年的分類任務(wù)改造希望幫到你。本文還有配套的精品資源點擊獲取