:從數(shù)據(jù)處理到遷移學習的完整流程)
1. 項目概述從一張圖片到識別結果到底發(fā)生了什么用 PyTorch 做圖片識別核心任務通常是圖像分類也就是我們輸入一張圖片模型輸出它屬于哪個類別。比如一張貓的照片經(jīng)過模型前向傳播得到貓 0.92、狗 0.05、鳥 0.02這樣的概率分布最后取最高概率對應的類別作為預測結果。很多初學者拿到這個需求時第一反應是調個模型跑一下但實際落地時你會發(fā)現(xiàn)整個流程是一條完整鏈路圖像讀取、預處理、數(shù)據(jù)加載、模型搭建、訓練迭代、驗證調參、固定權重、部署推理。任何一個環(huán)節(jié)處理不當訓練出來的模型精度都會差得離譜。這篇文章我會把這條鏈路每一個環(huán)節(jié)的實操細節(jié)都拆開講清楚包括一些你翻文檔翻不到的經(jīng)驗和坑。這篇文章適合誰看剛接觸 PyTorch、想快速跑通一個圖像分類項目的同學或者已經(jīng)能跑通官方 Demo 但換了自己的數(shù)據(jù)集就不知所措的朋友。我會用貓狗分類作為貫穿全文的示例任務因為它的數(shù)據(jù)獲取簡單、類別直觀、訓練收益明顯。文章里涉及的概念和代碼寫法換成任何其他分類任務都通用。2. 核心鏈路拆解圖片識別的五段流程2.1 圖像分類的本質圖像分類是計算機視覺最基礎的任務。它的數(shù)學本質是學習一個映射函數(shù) f輸入一張 H×W×C 的圖片張量輸出一個長度為 K 的向量K 是類別總數(shù)向量中的值經(jīng)過 Softmax 歸一化后表示每個類別的概率。在這里有一個很多新手沒想明白的點圖片輸入到網(wǎng)絡里其實是一堆數(shù)字。比如一張 224×224 的 RGB 圖片輸入張量的形狀是 [3, 224, 224]取值范圍在 0 到 1 之間。卷積神經(jīng)網(wǎng)絡在做的事情本質上是從這些數(shù)字中逐層提取特征——淺層卷積提取邊緣、顏色、紋理深層卷積提取眼睛、耳朵、輪子這種語義級別的部件。最后一層全連接將這些特征映射到類別空間。所以整個識別流程的五個環(huán)節(jié)是圖像預處理 → 數(shù)據(jù)加載 → 模型前向傳播 → 損失計算與反向傳播 → 推理預測。環(huán)境準備完成后后面每個環(huán)節(jié)我會單獨展開講。2.2 環(huán)境準備與硬件選型先說你需要的環(huán)境這里給出一個我用下來比較穩(wěn)定的版本組合PyTorch 2.x CUDA 11.8 或 12.1 Python 3.9~3.11 torchvision 配套版本。嚴格來說我們不建議追求最新版本穩(wěn)定性和兼容性更重要很多 GPU 驅動對最新的 CUDA 版本支持滯后跑起來反而出各種詭異報錯。有兩點經(jīng)驗說明一下。第一如果你用的是 NVIDIA 顯卡安裝 PyTorch 之前先去官網(wǎng)根據(jù)你的 CUDA 版本選擇對應的安裝命令不要直接用 pip install torch 默認裝 CPU 版本不然訓練速度會慢幾十倍。第二如果沒有獨立顯卡也可以先跑通整個流程只是訓練時間會長很多本文所有代碼在 CPU 上都能運行只不過我們建議把 batch size 調小。3. 數(shù)據(jù)處理與加載比模型更影響精度的地方3.1 數(shù)據(jù)目錄如何組織PyTorch 的 torchvision.datasets.ImageFolder 是處理分類任務最省事的工具前提是你的數(shù)據(jù)目錄按類別組織成這樣的結構data/ train/ cat/ cat_001.jpg cat_002.jpg dog/ dog_001.jpg dog_002.jpg val/ cat/ cat_001.jpg dog/ dog_001.jpgtrain 和 test 文件夾各自包含類別子文件夾里面的圖片按類別放置。用這種方式組織數(shù)據(jù)ImageFolder 會自動讀取所有圖片并根據(jù)子文件夾名稱生成類別索引我們完全可以用自己真實的數(shù)據(jù)目錄來替換這個示例結構。一個要注意的點訓練集和驗證集必須分開。我見過不少朋友把所有圖片放在一個大文件夾里直接訓練這樣驗證集失去了意義。因為模型見過了所有數(shù)據(jù)驗證集精度虛高換到真實場景性能崩盤。劃分比例一般是 82 或 91。3.2 Dataset 與 DataLoader 的關鍵配置有了目錄結構接下來就是標準的數(shù)據(jù)加載代碼。這段代碼幾乎是所有 PyTorch 圖像分類項目的骨架你需要根據(jù)自己的數(shù)據(jù)集特點去調整里面的參數(shù)。import torch from torch.utils.data import DataLoader from torchvision import datasets, transforms # 訓練集預處理 train_transform transforms.Compose([ transforms.Resize((256, 256)), # 先縮放 transforms.RandomCrop(224), # 再隨機裁剪 transforms.RandomHorizontalFlip(), # 隨機水平翻轉 transforms.ToTensor(), # 轉為張量并歸一化到 [0,1] transforms.Normalize( mean[0.485, 0.456, 0.406], # ImageNet 均值 std[0.229, 0.224, 0.225] # ImageNet 標準差 ) ]) # 驗證集預處理只做縮放和歸一化不做數(shù)據(jù)增強 val_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_dataset datasets.ImageFolder(data/train, transformtrain_transform) val_dataset datasets.ImageFolder(data/val, transformval_transform) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size64, shuffleFalse, num_workers4, pin_memoryTrue)這里面的重點在于為什么訓練要用縮放 隨機裁剪而驗證集直接縮放。訓練集需要數(shù)據(jù)增強來增加樣本多樣性隨機裁剪相當于在 256×256 的圖片中隨機截取 224×224 的區(qū)域相當于一個樣本生成了很多不同的子樣本能有效緩解過擬合。而驗證集的目的是評估模型的真實表現(xiàn)所以只用固定縮放保證一致性。另一個容易忽略的配置是 shuffle 參數(shù)。訓練集必須 shuffleTrue因為如果不打亂每個 batch 里都是同一類別的圖片梯度下降的方向會來回震蕩收斂速度慢且不穩(wěn)定。驗證集則不需要 shuffle。關于 Normalize 的 mean 和 std這一組值來自 ImageNet 數(shù)據(jù)集的統(tǒng)計結果。我們做遷移學習時通常沿用這組值因為它和預訓練模型的輸入分布是匹配的。如果從頭訓練自己的模型理論上應該重新統(tǒng)計自己數(shù)據(jù)集的均值和標準差但實際操作中大多數(shù)情況下用了 ImageNet 的數(shù)值也不會差太多。3.3 數(shù)據(jù)增強做多少才算合適數(shù)據(jù)增強是把雙刃劍。太少模型容易過擬合太多會給訓練制造不必要的難度。以貓狗分類為例我常用的增強組合按優(yōu)先級排列隨機水平翻轉必選不影響語義隨機裁剪必選模擬尺度變化隨機顏色抖動可選增強光照魯棒性隨機旋轉最多 ±15 度太大會破壞圖像結構如果數(shù)據(jù)集只有幾千張圖片我建議增強力度適當加強把 RandomResizedCrop 的尺度范圍設為 (0.5, 1.0)讓模型學會從不同尺度下識別目標。如果數(shù)據(jù)集有十萬級別增強可以適度保守重點放在防止過擬合上。我做項目時驗證過一個經(jīng)驗當訓練集樣本量少于 5000 時數(shù)據(jù)增強的收益是顯著的分類精度可以提升 2 到 5 個百分點。但如果你用的樣本量已經(jīng)很大或者數(shù)據(jù)本身質量很高盲目加增強反而可能導致精度不升反降。4. 模型如何選手寫 CNN 還是遷移學習4.1 從零搭建一個可用的 CNN如果從頭搭建一個卷積神經(jīng)網(wǎng)絡核心是感受野的設計。簡單來說每一層卷積核覆蓋的區(qū)域大小決定了網(wǎng)絡能看到多大范圍的圖像信息。一個經(jīng)典的小型 CNN 結構如下import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self, num_classes2): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 32, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(64, 128, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), ) self.classifier nn.Sequential( nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(128, 128), nn.ReLU(inplaceTrue), nn.Dropout(0.5), nn.Linear(128, num_classes) ) def forward(self, x): x self.features(x) x self.classifier(x) return x這個模型的思路是輸入 3×224×224 圖片經(jīng)過三層卷積 池化后特征圖變成 128×28×28然后通過 AdaptiveAvgPool2d 把特征圖壓縮成 1×1再接全連接分類。三層卷積提取不同層次的特征池化逐步降低分辨率擴大感受野。從零搭建 CNN 的優(yōu)勢是不依賴外部權重部署時自由度很大訓練過程完全可控。缺點是精度天花板明顯對于復雜任務很難超過預訓練模型。所以我的建議是如果數(shù)據(jù)集很小但圖像特征明顯比如純色背景下區(qū)分圓形和方形自建 CNN 足夠用如果圖像復雜、類別相近優(yōu)先考慮遷移學習。4.2 遷移學習用預訓練模型大幅提升精度遷移學習的邏輯和自建模型完全不同。它的核心是復用在大規(guī)模數(shù)據(jù)集上訓練好的模型參數(shù)——比如在 ImageNet 上已經(jīng)學好了邊緣、紋理、形狀等通用特征——然后在自己數(shù)據(jù)集上微調。這就像你不需要重新學習認識顏色和輪廓只需要學習貓和狗的區(qū)別就行。使用 torchvision 的預訓練模型最常用的幾個選項ResNet18 / ResNet50精度和速度最均衡的選擇EfficientNet-B0計算量更小精度相當MobileNetV3適合移動端部署場景VGG16結構簡單但參數(shù)量大訓練慢實用性下降以 ResNet18 為例遷移學習的標準寫法是凍結特征層參數(shù)替換最后一層全連接from torchvision import models model models.resnet18(weightsmodels.ResNet18_Weights.DEFAULT) # 凍結所有層的參數(shù) for param in model.parameters(): param.requires_grad False # 獲取最后一層全連接的輸入維度 num_features model.fc.in_features # 替換成自己的分類器 model.fc nn.Sequential( nn.Linear(num_features, 256), nn.ReLU(inplaceTrue), nn.Dropout(0.5), nn.Linear(256, 2) )這段代碼里最關鍵的邏輯是 requires_grad False。凍結特征層參數(shù)后反向傳播的梯度不會更新這些層的權重模型只會訓練我們新加的全連接層。這樣有幾個好處一是訓練速度大幅提升因為不需要計算大量參數(shù)的梯度二是可以有效防止數(shù)據(jù)集過小導致的過擬合避免破壞預訓練模型已經(jīng)學好的通用特征表示。微調有兩種常用策略一是上面這種只訓練分類頭的方案適合數(shù)據(jù)集幾百到幾千張的情況二是全量微調也就是所有層都放開訓練適合數(shù)據(jù)集規(guī)模較大時有足夠的樣本來刷新預訓練知識。實際項目中我建議先用凍結方式訓練幾個 epoch 讓分類頭收斂然后解凍最后幾層做微調這樣精度更高也穩(wěn)定。4.3 損失函數(shù)和優(yōu)化器的選擇邏輯分類任務的標準配置是交叉熵損失函數(shù)加 SGD 或 Adam。這里的交叉熵損失函數(shù)值得多說一句PyTorch 的 nn.CrossEntropyLoss 其實做了兩件事——對模型的輸出做 Softmax 歸一化然后計算負對數(shù)似然損失。所以模型最后一層不需要再接 Softmax直接輸出原始 logits 就好否則等于做了兩次歸一化精度會受影響。criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.fc.parameters(), lr0.001)優(yōu)化器方面Adam 適合快速迭代和調參對學習率的敏感度低一些。SGD 加 Momentum 需要更精細的學習率調整但收斂后的精度往往更高。我在實際訓練中習慣先用 Adam 跑通流程之后如果有優(yōu)化需求再切換到帶學習率衰減的 SGD。這里有個新手常踩的坑如果只訓練分類頭優(yōu)化的參數(shù)只有 model.fc.parameters()因為被凍結的層沒有梯度。如果全量微調則要傳入 model.parameters()。如果混用凍結層和可訓練層需要確保優(yōu)化器中傳入了正確的參數(shù)組否則會出現(xiàn)模型一點也不學習的困惑現(xiàn)象。4.4 設備準備讓代碼自動跑在 GPU 或 CPU 上模型訓練前要先確定計算設備。建議直接寫成下面這種自適應方式避免在 GPU 機器上調試過的代碼挪到 CPU 機器上直接報錯device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device)model.to(device) 這一步很多人會漏掉導致前向傳播時模型還在 CPU 上數(shù)據(jù)在 GPU 上報錯說 device mismatch。還有一點需要特別注意輸入圖片張量也要調用 .to(device) 移動設備我在初學階段經(jīng)常因為忘記把 data 移動到 GPU 上而報錯這種問題排查起來很浪費時間。5. 訓練循環(huán)與模型驗證的完整實現(xiàn)5.1 標準訓練循環(huán)代碼訓練循環(huán)是整個項目最核心的代碼段。我把一個完整的訓練驗證循環(huán)貼出來這份代碼可以直接用以下方式運行然后根據(jù)打印的日志判斷訓練效果。def train_one_epoch(model, train_loader, criterion, optimizer, device): model.train() running_loss 0.0 correct 0 total 0 for inputs, labels in train_loader: inputs, labels inputs.to(device), labels.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * inputs.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() epoch_loss running_loss / total epoch_acc correct / total return epoch_loss, epoch_acc def validate(model, val_loader, criterion, device): model.eval() running_loss 0.0 correct 0 total 0 with torch.no_grad(): for inputs, labels in val_loader: inputs, labels inputs.to(device), labels.to(device) outputs model(inputs) loss criterion(outputs, labels) running_loss loss.item() * inputs.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() epoch_loss running_loss / total epoch_acc correct / total return epoch_loss, epoch_acc一個容易犯錯的地方是寫法上訓練時必須調用 model.train()驗證時必須調用 model.eval()。這兩個模式的區(qū)別在于 BatchNorm 和 Dropout 的行為train 模式下 Dropout 會隨機丟棄神經(jīng)元BatchNorm 會使用 batch 內的統(tǒng)計量eval 模式下 Dropout 關閉BatchNorm 使用訓練階段累積的全局統(tǒng)計量。如果驗證時忘了切換到 eval 模式驗證精度會異常波動并且吃進隨機性。另外驗證階段必須包裹 torch.no_grad()它有兩個用處一是停止梯度追蹤節(jié)省顯存二是避免驗證過程意外修改模型參數(shù)。5.2 模型保存策略只保留最優(yōu)權重訓練過程中的模型保存也是有講究的。不經(jīng)思考地只在最后保存一次模型如果你中途過擬合了最后保存的就是過擬合的模型。我建議采用驗證精度最高時保存的策略best_acc 0.0 for epoch in range(num_epochs): train_loss, train_acc train_one_epoch(...) val_loss, val_acc validate(...) print(fEpoch {epoch1}: train_loss{train_loss:.4f}, train_acc{train_acc:.4f}, fval_loss{val_loss:.4f}, val_acc{val_acc:.4f}) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_model.pth) print(f - 保存最優(yōu)模型驗證精度 {best_acc:.4f})torch.save 建議只保存 model.state_dict()它只保存模型權重參數(shù)不攜帶完整的模型結構。原因有兩個一是文件體積更小二是加載時更靈活。你在 state_dict 之外還需要單獨記錄類別映射關系否則推理時只有數(shù)字類別沒有類別名稱。這個我后面會講。5.3 學習率調整和 early stopping 的實際價值訓練過程中學習率設置是一個常見痛點。學習率太大模型震蕩不收斂太小則訓練半天損失紋絲不動。我喜歡用 ReduceLROnPlateau 做驗證精度的自適應調整scheduler torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, modemin, factor0.5, patience3 ) for epoch in range(num_epochs): train_loss, train_acc train_one_epoch(...) val_loss, val_acc validate(...) scheduler.step(val_loss)這里傳入的 modemin 表示當 val_loss 不再下降時把學習率乘以 0.5。patience3 意味著連續(xù) 3 個 epoch 驗證損失沒有明顯下降才觸發(fā)衰減。學習率衰減的原理是訓練初期需要較大的步長快速靠近最優(yōu)點但接近最優(yōu)點后大步長會在最優(yōu)值附近震蕩此時縮小步長可以讓損失繼續(xù)下降。early stopping 指的是當驗證損失連續(xù)多個 epoch 不再下降時提前終止訓練。這個和保存最優(yōu)模型配合起來用能有效避免過擬合。但我實際使用中有一個心得early stopping 的 patience 設置要根據(jù)數(shù)據(jù)集大小靈活調整小數(shù)據(jù)集設 5 個 epoch 比較合理大數(shù)據(jù)集可能需要 10 個甚至更多 epoch 的耐心否則容易在驗證損失的局部極小值處提前停止。6. 推理階段加載模型識別單張新圖片6.1 完整的推理代碼訓練完成后真正到識別圖片這個動作時代碼要包含三個步驟加載模型權重、對輸入圖片做和訓練時完全一致的預處理、前向傳播得到預測結果。這份推理代碼我是按照生產(chǎn)環(huán)境的標準寫的可以直接用import torch from PIL import Image from torchvision import transforms def predict_image(image_path, model, device, class_names): # 1. 加載并預處理圖片 image Image.open(image_path).convert(RGB) # 2. 預處理與驗證集的預處理完全一致 preprocessing transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) input_tensor preprocessing(image).unsqueeze(0) # 形狀 [1, 3, 224, 224] input_tensor input_tensor.to(device) # 3. 推理 model.eval() with torch.no_grad(): outputs model(input_tensor) probabilities torch.softmax(outputs, dim1) confidence, predicted_idx torch.max(probabilities, dim1) predicted_class class_names[predicted_idx.item()] return predicted_class, confidence.item()調用示例# 設定類別映射0 對應 cat1 對應 dog class_names [cat, dog] # 加載訓練好的最優(yōu)權重 model models.resnet18(weightsNone) num_features model.fc.in_features model.fc nn.Sequential(...) # 結構必須和訓練時完全一致 model.load_state_dict(torch.load(best_model.pth, map_locationdevice)) model model.to(device) result, confidence predict_image(test_photo.jpg, model, device, class_names) print(f識別結果{result}置信度{confidence:.2%})這里的關鍵坑在于 load_state_dict 前模型結構必須與訓練時完全一致。如果你在訓練代碼里改過 model.fc 的結構那么推理代碼里也要做一模一樣的修改否則加載權重時會報 size mismatch。為了避免這種問題我強烈建議把模型構建代碼寫在一個函數(shù)里訓練和推理都調用同一個函數(shù)不要二次復制。6.2 置信度識別結果的第二落點很多人做圖片識別只看預測類別但置信度信息其實更具價值。置信度是 Softmax 輸出的最大概率值表示為模型對該預測的自信程度。我在實際項目中總結的置信度決策規(guī)則置信度 0.9模型非常確定可以放心采納置信度在 0.6 到 0.9 之間模型有一定把握可以采納但要提示用戶核實置信度低于 0.6建議返回無法確定因為多類別中每個類別的概率都比較分散強取最大值很可能是錯誤判斷一個常見場景是模型對貓的置信度是 0.45對狗的置信度是 0.42此時取貓作為結果并不穩(wěn)妥。這種低置信度通常意味著圖片里的目標不明顯、遮擋嚴重或包含訓練數(shù)據(jù)中未出現(xiàn)的背景干擾。如果業(yè)務場景對誤判容忍度低建議加一個置信度閾值判斷。6.3 推理性能優(yōu)化經(jīng)驗如果識別功能要集成到實際應用中推理速度和延遲就很重要了。幾個實用的優(yōu)化方向把整個模型用 model.eval() 切換去掉 Dropout 的隨機性對輸入圖片做 batch 推理把多張圖片合成一個 batch 一次前向傳播比單張循環(huán)快很多倍使用 torch.no_grad() 防止構建計算圖如果需要更高吞吐量可以考慮用 ONNX 導出模型加速但在精度要求高時要注意導出和推理的精度差異我參與過的一個圖片識別的技術 Demo用 ResNet18 在 GPU 上單張推理延遲只有幾毫秒完全滿足實時應用的需求。而如果換用 CPU 推理延遲會上升到幾十毫秒這種差距會讓一些對時延敏感的應用場景做出不同的技術決策。7. 訓練過程中遇到的經(jīng)典問題排查實錄7.1 損失不下降我接觸的初學者項目里損失不下降或下降極慢是最常見的問題。排查順序應該按優(yōu)先級這樣走首先是學習率。學習率太小相當于一個人走路步子邁得極小半天走不出原地。把學習率從默認的 0.001 調到 0.01 或 0.1 測試一下。如果損失反而變成 NaN說明學習率又太大了中間找平衡點。我常用的粗暴方式是 0.1、0.01、0.001 各試 20 個 epoch看哪個收斂快且穩(wěn)定。其次是數(shù)據(jù)和標簽是否對齊。檢查 dataloader 返回的圖片和標簽是否匹配我遇到過圖片讀取順序和標簽順序錯位的問題導致模型一直在學錯誤對應關系??梢栽?DataLoader 里加調試代碼打印一批圖片對應的標簽確認。還有一個隱蔽的問題預處理時沒有做歸一化或歸一化錯誤。如果圖片像素值還是 0~255 范圍模型第一層卷積的輸入范圍就不對特別是使用遷移學習模型時預訓練權重期望的是 0~1 歸一化后帶標準差的輸入輸入范圍錯誤會讓特征提取完全崩潰。7.2 顯存溢出顯存溢出Out of Memory在小顯存機器上訓練大模型時幾乎必現(xiàn)。最有效的排查思路是降低 batch size從 32 降到 16 甚至 8。因為顯存占用主要是每個 batch 的中間特征圖和梯度batch size 直接線性影響顯存。另一個思路是減少圖片輸入尺寸224 降到 160 或 128這對顯存的影響是平方級別的。還有一招是使用梯度累積也就是用多個小 batch 模擬大 batch 的效果。比如顯存只夠跑 batch size 8但你需要 batch size 32那就跑 4 個 batch 后累積梯度再更新一次參數(shù)。7.3 過擬合怎么判斷判斷過擬合不能只看訓練集精度因為訓練集精度沖到 100% 也很正常。核心指標是訓練和驗證差距當訓練精度 95% 以上而驗證精度停留在 75% 左右且驗證損失在下降后又開始回升這就是典型過擬合。緩解手段按優(yōu)先級排序數(shù)據(jù)增強加碼、加 Dropout 比例、縮小模型規(guī)模、提前停止。不要上來就換更大的模型那樣只會過擬合得更嚴重。我在做項目時首先看的是數(shù)據(jù)增強有沒有做夠數(shù)據(jù)集的多樣性是否覆蓋了應用場景的分布。7.4 類別不平衡問題實際項目中類別不平衡很常見。比如你做一個缺陷檢測99% 的圖片是合格品只有 1% 是有缺陷的。模型不需要學任何東西全預測合格就有 99% 精度但這種模型完全沒用。解決辦法一是加權損失函數(shù)PyTorch 的 CrossEntropyLoss 支持 weight 參數(shù)class_weights torch.tensor([1.0, 10.0]) # 少數(shù)類權重設為大的值 criterion nn.CrossEntropyLoss(weightclass_weights.to(device))解決辦法二是采樣策略對少數(shù)類過采樣、對多數(shù)類欠采樣。我自己更傾向于先用加權損失它簡單直接改動量小實用性高。7.5 亂加預訓練權重導致的報錯用遷移學習模型時遇到 state_dict 加載報錯一般是新舊模型結構不一致。最常發(fā)生的情況是你用了一個 torchvision 模型然后手動修改了 model.fc但加載預訓練權重時服務端期望的還是原始完整結構。解決辦法是設置 weightsNone 構造模型后再加載自己的權重同時確保模型結構完全一致。8. 實操總結與擴展方向用 PyTorch 識別圖片這個需求看起來簡單但把一個模型真正跑通、跑好中間的信息量其實非常大。我最后分享幾個經(jīng)驗數(shù)據(jù)永遠是第一位的。模型結構再先進數(shù)據(jù)沒準備好效果也出不來。你花在數(shù)據(jù)清洗、數(shù)據(jù)增強上的時間永遠比調模型結構參數(shù)更值得。遷移學習是現(xiàn)在絕大多數(shù)場景的首選方案。自己從頭訓練一個卷積神經(jīng)網(wǎng)絡除非你有幾十萬上百萬級的標注數(shù)據(jù)否則很難超過一個微調好的預訓練模型。用預訓練模型不是偷懶是把前人幾十年的經(jīng)驗拿過來用這是聰明人的選擇。訓練日志一定要完整記錄。損失值、精度、學習率、模型在哪個 epoch 保存的這些信息看起來不起眼但在你回溯問題、復現(xiàn)實驗時能省下大把時間。我自己現(xiàn)在每跑一個實驗都會把超參數(shù)和日志文件一起歸檔。做完一個基礎分類項目后你可以順這個方向繼續(xù)深入一是多標簽分類一張圖中同時識別多個物體二是目標檢測不只是分類而是框出物體的位置三是語義分割對每個像素做類別預測。這三個方向都是在圖像識別基礎上的自然延伸掌握了本文的基礎流程后你會發(fā)現(xiàn)它們很多思維方式是一脈相承的。希望這套完整的流程和代碼能夠幫你把第一個 PyTorch 圖像識別項目順利跑通。在實際操作中遇到問題歡迎把你的報錯信息和實驗現(xiàn)象記錄下來大家一起交流排查的思路。