前陣子有個(gè)朋友問(wèn)我想用PyTorch做一個(gè)圖像識(shí)別的小項(xiàng)目但他翻了一堆教程要么只跑通了單張圖片的預(yù)測(cè)演示要么一上來(lái)就丟一個(gè)幾百層的模型讓人直接復(fù)制。我聽(tīng)完就覺(jué)得這類(lèi)教程最大的問(wèn)題是把“能跑”和“能落地”混為一談了。所以這篇文章我不想復(fù)述官方文檔而是把我實(shí)際做過(guò)的圖像識(shí)別完整流程拆開(kāi)講從任務(wù)定義、數(shù)據(jù)準(zhǔn)備、網(wǎng)絡(luò)選型到訓(xùn)練調(diào)參、模型評(píng)估和部署前的最后一步。文章適合正在學(xué)PyTorch、打算拿它做圖像識(shí)別的初學(xué)者也適合已經(jīng)跑過(guò)一些Demo、但不確定下一步怎么走的中級(jí)使用者。我盡量把每一步背后的“為什么”也講清楚而不是只給代碼。1. 動(dòng)手前先想清楚你的圖像識(shí)別任務(wù)到底是什么1.1 先把“圖像識(shí)別”翻譯成計(jì)算機(jī)能解決的問(wèn)題很多人上來(lái)就搜“PyTorch圖像識(shí)別”這個(gè)詞其實(shí)太寬泛了。圖像識(shí)別在工程上至少會(huì)分化成圖像分類(lèi)、目標(biāo)檢測(cè)、圖像分割、圖像檢索等好幾個(gè)分支技術(shù)方案完全不一樣。比如目標(biāo)檢測(cè)要輸出“圖中某個(gè)區(qū)域內(nèi)有哪類(lèi)目標(biāo)”用的是邊界框回歸加分類(lèi)圖像分割則要像素級(jí)標(biāo)記需要U-Net這類(lèi)結(jié)構(gòu)。如果你的需求只是“這張圖是什么”那就對(duì)應(yīng)單標(biāo)簽多分類(lèi)問(wèn)題這是入門(mén)最友好的一個(gè)方向。我建議動(dòng)手之前先把任務(wù)參數(shù)寫(xiě)死類(lèi)別數(shù)量、每類(lèi)大概多少樣本、輸入圖片的尺寸、圖片里的目標(biāo)是否足夠顯眼。另一個(gè)容易被忽略的點(diǎn)是數(shù)據(jù)來(lái)源。如果你只有幾十張圖幾百?gòu)垐D就要考慮遷移學(xué)習(xí)如果每類(lèi)有幾千張甚至更多從頭訓(xùn)練一個(gè)網(wǎng)絡(luò)才說(shuō)得通。把這些約束列清楚后面每一步選擇都有據(jù)可依而不是憑感覺(jué)調(diào)參。1.2 為什么現(xiàn)在大多數(shù)學(xué)圖像識(shí)別的人都選PyTorchPyTorch在圖像識(shí)別領(lǐng)域流行核心原因是它的動(dòng)態(tài)圖機(jī)制。簡(jiǎn)單說(shuō)你用代碼寫(xiě)著寫(xiě)著就能隨時(shí)打印中間張量的shape可以在中間層打斷點(diǎn)看數(shù)據(jù)流這對(duì)調(diào)試神經(jīng)網(wǎng)絡(luò)來(lái)說(shuō)太重要了。相比之下靜態(tài)圖模式相當(dāng)于先整體編譯再執(zhí)行出了問(wèn)題排錯(cuò)成本高很多。做圖像識(shí)別的時(shí)候張量維度稍微搞錯(cuò)一點(diǎn)就報(bào)錯(cuò)能直接print出來(lái)看每一層輸出排錯(cuò)速度完全不一樣。另外PyTorch的生態(tài)對(duì)圖像任務(wù)足夠友好。torchvision集成了大量預(yù)訓(xùn)練模型、常用數(shù)據(jù)集和數(shù)據(jù)增強(qiáng)工具DataLoader配合Dataset能很舒服地處理圖片數(shù)據(jù)。就算后續(xù)要部署到線(xiàn)上PyTorch也可以導(dǎo)出成ONNX格式給其他推理引擎消費(fèi)。這些細(xì)節(jié)讓我在實(shí)際項(xiàng)目里很少需要換框架。1.3 數(shù)據(jù)量、算力、精度目標(biāo)決定你的方案選網(wǎng)絡(luò)結(jié)構(gòu)不能只看哪個(gè)模型刷榜刷得高要看你的資源和約束。我的經(jīng)驗(yàn)可以分成三類(lèi)情況供你參考。場(chǎng)景數(shù)據(jù)規(guī)模算力情況推薦方案快速驗(yàn)證每類(lèi)幾百?gòu)垎螇K中端GPU預(yù)訓(xùn)練ResNet18 / MobileNet 只訓(xùn)練分類(lèi)頭常規(guī)項(xiàng)目每類(lèi)幾千張單塊好點(diǎn)GPU預(yù)訓(xùn)練ResNet50微調(diào)最后幾層 數(shù)據(jù)增強(qiáng)從零訓(xùn)練研究每類(lèi)上萬(wàn)張多卡或長(zhǎng)時(shí)間訓(xùn)練自己搭CNN或標(biāo)準(zhǔn)ResNet從頭訓(xùn)這個(gè)決策邏輯很直接遷移學(xué)習(xí)是用別人的特征提取能力來(lái)彌補(bǔ)自己數(shù)據(jù)不足的短板相當(dāng)于一個(gè)經(jīng)驗(yàn)豐富的師傅先把“看邊緣、看紋理、看形狀”的基本功練好了你只需要教它識(shí)別你的特定類(lèi)別。如果你數(shù)據(jù)量夠大從頭訓(xùn)練才可能超越預(yù)訓(xùn)練權(quán)重帶來(lái)的上限。2. 環(huán)境與數(shù)據(jù)準(zhǔn)備把底子打好2.1 環(huán)境安裝里你大概率會(huì)遇到的版本問(wèn)題我見(jiàn)過(guò)太多人卡在環(huán)境上不是PyTorch裝不上而是裝完之后運(yùn)行時(shí)才發(fā)現(xiàn)版本對(duì)不上。最穩(wěn)的流程是先創(chuàng)建一個(gè)獨(dú)立的虛擬環(huán)境避免和系統(tǒng)Python或者其他項(xiàng)目互相污染。如果你用conda可以這樣操作conda create -n classify python3.10 conda activate classify然后根據(jù)你是否使用GPU來(lái)安裝對(duì)應(yīng)版本。我的建議是無(wú)論你有沒(méi)有獨(dú)顯都先在終端里跑下面這些命令確認(rèn)一下python -c import torch; print(torch.__version__) python -c import torch; print(torch.cuda.is_available())如果第二行輸出的是True那GPU環(huán)境就通了。很多人明明按教程安裝但torch.cuda.is_available()依然是False最常見(jiàn)的原因是安裝時(shí)選了CPU版本或者驅(qū)動(dòng)版本太老。別急著重裝系統(tǒng)先查驅(qū)動(dòng)支持再去PyTorch官網(wǎng)選和你的CUDA能力匹配的安裝命令這樣最省時(shí)間。我踩過(guò)最冤的一個(gè)坑是conda自動(dòng)匹配了CPU版torch結(jié)果代碼跑了一天發(fā)現(xiàn)根本沒(méi)用GPU浪費(fèi)了大量時(shí)間。2.2 用文件夾結(jié)構(gòu)組織圖片比寫(xiě)幾百行代碼還管用數(shù)據(jù)準(zhǔn)備階段我最推薦的方式是把圖片按類(lèi)別放在文件夾里訓(xùn)練集、驗(yàn)證集、測(cè)試集分開(kāi)目錄管理。比如data/ train/ cat/ 圖片... dog/ 圖片... bird/ 圖片... val/ cat/ dog/ bird/ test/ cat/ dog/ bird/這種結(jié)構(gòu)的好處是torchvision.datasets.ImageFolder可以直接讀取它會(huì)自動(dòng)把子文件夾名稱(chēng)映射成類(lèi)別標(biāo)簽連標(biāo)簽列表都不用自己寫(xiě)。如果數(shù)據(jù)是一個(gè)CSV或者數(shù)據(jù)庫(kù)里存著路徑和標(biāo)簽?zāi)蔷托枰^承torch.utils.data.Dataset自定義讀取邏輯但多數(shù)入門(mén)項(xiàng)目用ImageFolder就夠了。from torchvision import datasets, transforms train_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) train_data datasets.ImageFolder(data/train, transformtrain_transform) print(train_data.classes) # [bird, cat, dog]注意Normalize里的均值和標(biāo)準(zhǔn)差。如果你用的是ImageNet預(yù)訓(xùn)練權(quán)重這三組數(shù)值幾乎是固定的如果你從頭訓(xùn)練自己的網(wǎng)絡(luò)可以用全量訓(xùn)練數(shù)據(jù)估算均值和標(biāo)準(zhǔn)差效果更貼合你的數(shù)據(jù)分布。但很多新手最容易犯的錯(cuò)誤就是對(duì)預(yù)訓(xùn)練模型用了和目標(biāo)數(shù)據(jù)集完全不一致的歸一化參數(shù)導(dǎo)致收斂困難。2.3 DataLoader參數(shù)逐個(gè)拆解batch_size、shuffle、num_workersDataLoader是連接數(shù)據(jù)集和訓(xùn)練循環(huán)之間的橋梁幾個(gè)參數(shù)直接影響訓(xùn)練速度和效果。from torch.utils.data import DataLoader train_loader DataLoader( train_data, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue ) val_loader DataLoader( val_data, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue )batch_size每次迭代喂給GPU的圖片數(shù)量。太大可能顯存溢出太小會(huì)讓梯度震蕩訓(xùn)練不穩(wěn)定。一般先在32或64里選OOM就減半驗(yàn)證集準(zhǔn)確率遲遲上不去就試著加大。shuffle訓(xùn)練集必須置為T(mén)rue否則每一輪epoch內(nèi)樣本順序不變模型可能學(xué)會(huì)記憶順序而非特征驗(yàn)證集和測(cè)試集必須置為False因?yàn)樵u(píng)估結(jié)果不應(yīng)該受打亂樣本順序影響。num_workers數(shù)據(jù)加載用了多少個(gè)進(jìn)程。這個(gè)數(shù)值越高CPU讀圖和預(yù)處理越快但不是越大越好Windows上經(jīng)常設(shè)置超過(guò)4會(huì)報(bào)錯(cuò)。pin_memory設(shè)為T(mén)rue會(huì)鎖頁(yè)內(nèi)存把數(shù)據(jù)從CPU傳輸?shù)紾PU的耗時(shí)降低。這個(gè)參數(shù)在數(shù)據(jù)量大的時(shí)候提升明顯代價(jià)是占用一部分內(nèi)存。3. 網(wǎng)絡(luò)結(jié)構(gòu)設(shè)計(jì)先跑通一個(gè)簡(jiǎn)單的CNN再談ResNet3.1 卷積網(wǎng)絡(luò)是怎么“看”圖像的如果你第一次接觸卷積神經(jīng)網(wǎng)絡(luò)可以把它理解成一臺(tái)滑動(dòng)的掃描儀。卷積核就是在圖片上從左到右、從上到下地滑動(dòng)每次都只觀察一個(gè)局部區(qū)域然后把局部區(qū)域的特征提取出來(lái)。相比全連接網(wǎng)絡(luò)卷積層最大的優(yōu)勢(shì)是參數(shù)共享和局部連接。同一組卷積核在整個(gè)畫(huà)幅上滑動(dòng)意味著它只需要學(xué)習(xí)“某種紋理長(zhǎng)什么樣”而不是“某個(gè)位置上有這種紋理怎么處理”參數(shù)量大幅下降也更不容易過(guò)擬合。每個(gè)卷積層輸出的特征圖尺寸可以用這樣一個(gè)基礎(chǔ)公式估算輸出邊長(zhǎng) (輸入邊長(zhǎng) - 卷積核大小 2 * padding) / stride 1。比如輸入224x224卷積核3x3padding為1stride為1輸出就是224x224。池化層則是對(duì)特征圖做降采樣可以類(lèi)比成把高清圖片縮小成縮略圖保留主要信息但丟掉冗余細(xì)節(jié)一定程度上也增強(qiáng)平移不變性。3.2 從零實(shí)現(xiàn)一個(gè)能跑的小型CNN我不建議初學(xué)者一上來(lái)就套大型網(wǎng)絡(luò)先手寫(xiě)一個(gè)輕量CNN把這些概念串起來(lái)后面換ResNet才心里有底。下面這個(gè)結(jié)構(gòu)是我常用的起步版輸入3通道224x224圖片輸出3類(lèi)import torch import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self, num_classes3): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 16, kernel_size3, padding1), nn.BatchNorm2d(16), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), # 224 - 112 nn.Conv2d(16, 32, kernel_size3, padding1), nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), # 112 - 56 nn.Conv2d(32, 64, kernel_size3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), # 56 - 28 ) self.classifier nn.Sequential( nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(64, num_classes) ) def forward(self, x): x self.features(x) x self.classifier(x) return x model SimpleCNN(num_classes3) dummy torch.randn(2, 3, 224, 224) out model(dummy) print(out.shape) # torch.Size([2, 3])我特意用了AdaptiveAvgPool2d它可以自適應(yīng)地把任意大小的特征圖池化成1x1這樣你的模型對(duì)輸入尺寸的寬容度會(huì)大很多換測(cè)試圖片尺寸時(shí)不用改全連接層的維度。這個(gè)細(xì)節(jié)是我后面做跨分辨率預(yù)測(cè)時(shí)經(jīng)常用到的。3.3 遷移學(xué)習(xí)把別人訓(xùn)練好的經(jīng)驗(yàn)借過(guò)來(lái)自己從零訓(xùn)練一個(gè)CNN常常需要大量數(shù)據(jù)和很長(zhǎng)時(shí)間實(shí)際項(xiàng)目中我更推薦直接在預(yù)訓(xùn)練模型基礎(chǔ)上做遷移學(xué)習(xí)。torchvision里加載預(yù)訓(xùn)練模型很直接from torchvision import models model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) model.fc nn.Linear(model.fc.in_features, num_classes) # 凍結(jié)大部分層只訓(xùn)練分類(lèi)頭 for param in model.parameters(): param.requires_grad False for param in model.fc.parameters(): param.requires_grad True凍結(jié)策略有幾種不同做法。最省資源和時(shí)間的是只訓(xùn)練最后一層分類(lèi)頭如果數(shù)據(jù)量中等可以嘗試解凍后面幾個(gè)block用較小的學(xué)習(xí)率訓(xùn)練如果你的數(shù)據(jù)已經(jīng)足夠多也可以不凍結(jié)任何層直接用預(yù)訓(xùn)練權(quán)重作為初始化全量微調(diào)。沒(méi)有絕對(duì)優(yōu)劣完全看你手里的數(shù)據(jù)量。我在實(shí)驗(yàn)中發(fā)現(xiàn)數(shù)據(jù)量只有每類(lèi)一千張左右時(shí)只訓(xùn)分類(lèi)頭反而比全量微調(diào)更穩(wěn)因?yàn)槿课⒄{(diào)容易在早期就過(guò)擬合。4. 訓(xùn)練細(xì)節(jié)準(zhǔn)確率提升的關(guān)鍵不在網(wǎng)絡(luò)而在訓(xùn)練4.1 數(shù)據(jù)增強(qiáng)用更少的圖片訓(xùn)練更強(qiáng)的模型很多人以為數(shù)據(jù)增強(qiáng)只是增加訓(xùn)練樣本的手段實(shí)際上它還有很強(qiáng)的正則化作用。隨機(jī)裁剪、水平翻轉(zhuǎn)、顏色抖動(dòng)這些變換相當(dāng)于告訴模型“目標(biāo)的類(lèi)別不應(yīng)該因?yàn)樗诋?huà)面中的位置、方向、亮度變化而改變”。最常見(jiàn)的組合是訓(xùn)練集做RandomResizedCrop加RandomHorizontalFlip加ColorJitter驗(yàn)證集只做縮放和歸一化。train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.8, 1.0)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) val_transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])注意驗(yàn)證集不要加隨機(jī)增強(qiáng)否則每次測(cè)試結(jié)果都會(huì)因?yàn)殡S機(jī)操作而抖動(dòng)沒(méi)法穩(wěn)定評(píng)估模型好壞。這是我最初最容易犯的一個(gè)原則性錯(cuò)誤一度懷疑自己的模型不穩(wěn)定。4.2 損失函數(shù)和優(yōu)化器的選擇邏輯圖像分類(lèi)最常用的損失函數(shù)是CrossEntropyLoss。它內(nèi)部已經(jīng)包含了Softmax的計(jì)算所以不需要再在網(wǎng)絡(luò)最后加一個(gè)Softmax層。這個(gè)損失函數(shù)對(duì)多分類(lèi)問(wèn)題的效果和穩(wěn)定性都很好我的建議是新手直接用它別折騰其他花式損失。優(yōu)化器方面SGD和Adam是兩大主流。Adam收斂快學(xué)習(xí)率可以設(shè)置得比較低尤其適合微調(diào)但SGD加上動(dòng)量在訓(xùn)練后期往往有更強(qiáng)的泛化能力。我自己的習(xí)慣是跑通流程先用Adam學(xué)習(xí)率0.001驗(yàn)證思路沒(méi)問(wèn)題后再換SGD試試能否進(jìn)一步提升準(zhǔn)確率。還有一個(gè)容易被忽視的參數(shù)是weight_decay它本質(zhì)上是給模型權(quán)重加了L2正則等于告訴模型“權(quán)重不要太極端”對(duì)控制過(guò)擬合有明顯幫助。代碼里就是一句optimizer torch.optim.Adam(model.parameters(), lr1e-3, weight_decay1e-4)4.3 學(xué)習(xí)率訓(xùn)練中最常改的一個(gè)參數(shù)學(xué)習(xí)率設(shè)置太大會(huì)導(dǎo)致?lián)p失發(fā)散太小會(huì)讓訓(xùn)練慢得讓人懷疑人生。我踩過(guò)很多次學(xué)習(xí)率過(guò)高導(dǎo)致loss變成NaN的坑特別是用預(yù)訓(xùn)練模型解鎖全量微調(diào)時(shí)整個(gè)網(wǎng)絡(luò)學(xué)習(xí)率一致上調(diào)很容易震碎已經(jīng)學(xué)好的權(quán)重。一個(gè)實(shí)用的經(jīng)驗(yàn)是新加的分類(lèi)頭用相對(duì)大的學(xué)習(xí)率預(yù)訓(xùn)練特征提取部分用較小的學(xué)習(xí)率。optimizer torch.optim.Adam([ {params: model.fc.parameters(), lr: 1e-3}, {params: model.layer4.parameters(), lr: 1e-4}, ], weight_decay1e-4)學(xué)習(xí)率調(diào)度同樣重要。我常用的是ReduceLROnPlateau也就是當(dāng)驗(yàn)證集準(zhǔn)確率不再上升時(shí)把學(xué)習(xí)率降為原來(lái)的十分之一。也可以用固定步長(zhǎng)的StepLR最簡(jiǎn)單最省心。訓(xùn)練早期如果加一個(gè)小的“warmup”比如前面幾個(gè)epoch把學(xué)習(xí)率從零逐漸線(xiàn)性升到目標(biāo)值可以在使用較大初始學(xué)習(xí)率時(shí)不那么容易被撐壞。4.4 一個(gè)規(guī)范化的訓(xùn)練循環(huán)別把驗(yàn)證邏輯寫(xiě)在同一步里訓(xùn)練循環(huán)看起來(lái)簡(jiǎn)單但很多人把訓(xùn)練和驗(yàn)證混在一起處理搞到最后連模型是哪個(gè)epoch保存的都不知道。我的習(xí)慣是每輪epoch做兩件事先在訓(xùn)練集上訓(xùn)練一遍再在驗(yàn)證集上完整評(píng)估一遍。兩個(gè)函數(shù)分開(kāi)寫(xiě)職責(zé)清晰。def train_one_epoch(model, loader, loss_fn, optimizer): model.train() total_loss, correct, total 0, 0, 0 for images, labels in loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss loss_fn(outputs, labels) loss.backward() optimizer.step() total_loss loss.item() correct (outputs.argmax(1) labels).sum().item() total labels.size(0) return total_loss / len(loader), correct / total def validate(model, loader, loss_fn): model.eval() total_loss, correct, total 0, 0, 0 with torch.no_grad(): for images, labels in loader: images, labels images.to(device), labels.to(device) outputs model(images) loss loss_fn(outputs, labels) total_loss loss.item() correct (outputs.argmax(1) labels).sum().item() total labels.size(0) return total_loss / len(loader), correct / total有兩個(gè)細(xì)節(jié)我特別想說(shuō)。第一model.train()和model.eval()必須在每個(gè)階段開(kāi)始前切換因?yàn)镈ropout和BatchNorm在訓(xùn)練和推理時(shí)的行為完全不同。第二驗(yàn)證階段一定用with torch.no_grad()包裹否則PyTorch會(huì)構(gòu)建推理圖顯存瞬間被吃掉而且白白增加計(jì)算量。5. 模型評(píng)估與落地從測(cè)試集到真實(shí)項(xiàng)目5.1 不要只看準(zhǔn)確率混淆矩陣和F1值能暴露更多問(wèn)題很多初學(xué)者的唯一指標(biāo)就是驗(yàn)證集準(zhǔn)確率。這個(gè)指標(biāo)在類(lèi)別均衡時(shí)還湊合但遇到類(lèi)別不均衡的數(shù)據(jù)集準(zhǔn)確率會(huì)非常騙人。舉個(gè)例子如果某個(gè)類(lèi)占了90%的樣本你的模型就算把所有圖片都預(yù)測(cè)成這個(gè)類(lèi)準(zhǔn)確率也有90%實(shí)際一點(diǎn)用都沒(méi)有。我建議每次訓(xùn)練完都看混淆矩陣把預(yù)測(cè)結(jié)果按行和列統(tǒng)計(jì)出來(lái)能直觀看到哪些類(lèi)別最容易互相混淆。然后計(jì)算每個(gè)類(lèi)別的精確率、召回率、F1值尤其關(guān)注數(shù)量少的類(lèi)別。用scikit-learn幾行代碼就能出來(lái)from sklearn.metrics import confusion_matrix, classification_report import numpy as np # preds和labels分別是所有預(yù)測(cè)結(jié)果和真實(shí)標(biāo)簽格式為L(zhǎng)ist[int] cm confusion_matrix(labels, preds) print(cm) print(classification_report(labels, preds, target_names[cat, dog, bird]))如果兩個(gè)類(lèi)別在混淆矩陣?yán)锟偸腔ハ嗾`判我一般會(huì)懷疑兩類(lèi)樣本本身存在較大的視覺(jué)相似性或者標(biāo)注質(zhì)量有問(wèn)題。這時(shí)候優(yōu)先檢查數(shù)據(jù)標(biāo)注其次是增加對(duì)應(yīng)類(lèi)別的樣本數(shù)量最后才考慮改網(wǎng)絡(luò)結(jié)構(gòu)。5.2 保存模型、加載模型和單張圖片預(yù)測(cè)訓(xùn)練結(jié)束后保存模型有兩種常見(jiàn)方式。只保存權(quán)重參數(shù)是最推薦的格式占空間小也方便后續(xù)換網(wǎng)絡(luò)結(jié)構(gòu)重新加載torch.save(model.state_dict(), model.pth)加載時(shí)要先構(gòu)造一個(gè)同樣結(jié)構(gòu)的模型再把權(quán)重放進(jìn)去model SimpleCNN(num_classes3) model.load_state_dict(torch.load(model.pth, map_locationcpu)) model.eval()推理階段和訓(xùn)練階段的寫(xiě)法完全不同。必須先切換到eval()模式再做一次和之前驗(yàn)證集相同的數(shù)據(jù)預(yù)處理然后調(diào)用模型from PIL import Image def predict_image(model, img_path, class_names): img Image.open(img_path).convert(RGB) img val_transform(img).unsqueeze(0) # 增加batch維度 model.eval() with torch.no_grad(): logits model(img) probs torch.softmax(logits, dim1) pred_idx logits.argmax(1).item() return class_names[pred_idx], probs[0][pred_idx].item() class_names train_data.classes pred, prob predict_image(model, test/cat_001.jpg, class_names) print(pred, prob)這里最容易翻車(chē)的點(diǎn)是忘了unsqueeze(0)。網(wǎng)絡(luò)默認(rèn)輸入的shape是[batch, channel, height, width]單張圖片讀進(jìn)來(lái)只有[channel, height, width]不加維度直接喂進(jìn)去就報(bào)維度錯(cuò)誤。另外一定要用softmax把輸出轉(zhuǎn)成概率分布否則你看到的就是一堆很難解釋的數(shù)值。5.3 把模型接到實(shí)際項(xiàng)目里導(dǎo)成ONNX交給推理引擎如果只是做實(shí)驗(yàn)?zāi)Ｐ屯Ｔ?pth文件就夠了。真要把識(shí)別功能集成到一個(gè)Web服務(wù)或者App里我建議導(dǎo)出ONNX格式。ONNX是一種開(kāi)放的模型交換格式很多推理引擎都能直接加載不依賴(lài)PyTorch環(huán)境本身。dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export(model, dummy_input, model.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}})這里設(shè)置dynamic_axes可以讓輸入的第一維變成動(dòng)態(tài)的這樣服務(wù)接口接收任意數(shù)量圖片都行。導(dǎo)出之后再用ONNX Runtime跑一遍輸入和輸出都應(yīng)該和PyTorch一致。能走到這一步一個(gè)圖像識(shí)別模型基本就具備了被實(shí)際服務(wù)調(diào)用的能力。6. 復(fù)現(xiàn)與效果分析我踩過(guò)的坑和最終對(duì)比6.1 最容易踩的幾個(gè)坑我按嚴(yán)重程度排了個(gè)序我梳理了自己在多個(gè)圖像識(shí)別項(xiàng)目里反復(fù)踩過(guò)的坑也是經(jīng)常在社區(qū)里看到新手們遇到的問(wèn)題從嚴(yán)重到輕微排列如下。問(wèn)題表現(xiàn)主要原因解決辦法Normalize參數(shù)用錯(cuò)模型完全無(wú)法收斂loss居高不下和預(yù)訓(xùn)練模型不匹配核對(duì)來(lái)源模型標(biāo)準(zhǔn)歸一化的均值和標(biāo)準(zhǔn)差訓(xùn)練時(shí)忘了model.train()使用帶Dropout的網(wǎng)絡(luò)時(shí)驗(yàn)證集很奇怪推理模式被用于訓(xùn)練每個(gè)階段切換train/eval驗(yàn)證集被隨機(jī)增強(qiáng)每次驗(yàn)證準(zhǔn)確率波動(dòng)大錯(cuò)誤復(fù)用了訓(xùn)練預(yù)處理驗(yàn)證集只用ResizeCenterCropNormalize隨機(jī)種子沒(méi)固定每次訓(xùn)練結(jié)果完全不同初始化、數(shù)據(jù)采樣順序隨機(jī)在代碼頂部設(shè)置random.seed、np.random.seed、手動(dòng)設(shè)置torch種子batch_size太大導(dǎo)致OOM訓(xùn)練中斷顯存溢出數(shù)據(jù)量和模型參數(shù)量超出顯存減半batch_size或者混合精度其中最容易被忽視的是隨機(jī)種子。如果你實(shí)驗(yàn)中發(fā)現(xiàn)連續(xù)兩次訓(xùn)練同一個(gè)模型準(zhǔn)確率差了1到2個(gè)百分點(diǎn)第一件事就是固定隨機(jī)種子。下面這段代碼我?guī)缀趺總€(gè)項(xiàng)目都復(fù)制import random import numpy as np import torch def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed)6.2 同樣一份數(shù)據(jù)下不同方案的效果差異有多大我給一個(gè)朋友做項(xiàng)目時(shí)用一份公開(kāi)的圖片分類(lèi)數(shù)據(jù)集做了四組對(duì)比實(shí)驗(yàn)每類(lèi)圖片數(shù)量在800到1200張之間。跑出來(lái)的趨勢(shì)很有參考價(jià)值方案訓(xùn)練輪數(shù)驗(yàn)證集準(zhǔn)確率訓(xùn)練耗時(shí)從頭訓(xùn)練SimpleCNN無(wú)增強(qiáng)4078%約20分鐘從頭訓(xùn)練SimpleCNN 數(shù)據(jù)增強(qiáng)4083%約25分鐘預(yù)訓(xùn)練ResNet18 只訓(xùn)分類(lèi)頭1586%約15分鐘預(yù)訓(xùn)練ResNet18 解凍layer4 增強(qiáng)2089%約20分鐘這個(gè)對(duì)比的結(jié)論很清楚數(shù)據(jù)增強(qiáng)帶來(lái)的提升比換網(wǎng)絡(luò)結(jié)構(gòu)更明顯預(yù)訓(xùn)練模型哪怕只訓(xùn)分類(lèi)頭也比從頭訓(xùn)練一個(gè)小網(wǎng)絡(luò)準(zhǔn)確率高出不少。而最后一種方案融合了增強(qiáng)和微調(diào)效果是最好的。實(shí)際解決業(yè)務(wù)問(wèn)題時(shí)我通常沿用這個(gè)遞進(jìn)思路而不是一上來(lái)就追求最復(fù)雜的大模型。6.3 從頭復(fù)現(xiàn)這套流程的操作順序把上面內(nèi)容收攏成一個(gè)可以照著做的清單我的建議順序是定義分類(lèi)任務(wù)確認(rèn)類(lèi)別數(shù)量和數(shù)據(jù)量。建立train/val/test目錄核對(duì)每個(gè)類(lèi)別的樣本數(shù)和圖片質(zhì)量。固定隨機(jī)種子搭建DataLoader和基礎(chǔ)訓(xùn)練循環(huán)先用一個(gè)非常小的網(wǎng)絡(luò)過(guò)一遍全流程。把網(wǎng)絡(luò)換成預(yù)訓(xùn)練ResNet18先只訓(xùn)分類(lèi)頭確保能收斂。加入數(shù)據(jù)增強(qiáng)和學(xué)習(xí)率調(diào)度提升最終準(zhǔn)確率。用混淆矩陣和F1檢查每個(gè)類(lèi)別的表現(xiàn)定位易混淆類(lèi)別。保存模型導(dǎo)出ONNX接進(jìn)實(shí)際應(yīng)用。這套順序的核心原則是先讓流程轉(zhuǎn)起來(lái)再逐步優(yōu)化。我見(jiàn)過(guò)太多人一上來(lái)就在網(wǎng)絡(luò)結(jié)構(gòu)上花大量時(shí)間結(jié)果連基礎(chǔ)的數(shù)據(jù)加載和訓(xùn)練循環(huán)都沒(méi)跑通最后反而被一個(gè)維度錯(cuò)誤卡了一整天。最后說(shuō)一個(gè)我自己的習(xí)慣每改一次數(shù)據(jù)預(yù)處理、網(wǎng)絡(luò)結(jié)構(gòu)或者損失函數(shù)我都會(huì)先固定隨機(jī)種子在小數(shù)據(jù)上跑幾個(gè)epoch確認(rèn)流程沒(méi)問(wèn)題再上全量數(shù)據(jù)。這個(gè)小習(xí)慣幫我過(guò)濾掉了大量“模型效果變差但原因卻在于流程本身出bug”的虛假實(shí)驗(yàn)結(jié)論。圖像識(shí)別這個(gè)方向看起來(lái)模型很炫真正決定項(xiàng)目成敗的往往都是這些訓(xùn)練流程里的基本功。