一視覺分類模板:從AlexNet到ViT訓(xùn)練與ONNX部署)
很多剛接觸視覺算法的同學(xué)都會遇到一個類似困境在文章或課程里把 AlexNet、VGG、ResNet、ViT 的結(jié)構(gòu)圖看得明明白白真正打開編輯器準(zhǔn)備訓(xùn)練時卻發(fā)現(xiàn)自己只有“網(wǎng)絡(luò)結(jié)構(gòu)代碼”沒有一套完整的數(shù)據(jù)加載、訓(xùn)練循環(huán)、評估和導(dǎo)出腳本。于是從 A 倉庫復(fù)制數(shù)據(jù)增強從 B 倉庫找訓(xùn)練參數(shù)再從 C 倉庫拼接模型調(diào)用最后訓(xùn)練到一半才發(fā)現(xiàn)學(xué)習(xí)率策略和優(yōu)化器根本不匹配。這不是個人能力問題而是當(dāng)前視覺入門教程普遍存在的空白結(jié)構(gòu)講解很多工程閉環(huán)太少??梢越o出一個明確判斷在圖像分類這個方向上從經(jīng)典 CNN 到 Transformer訓(xùn)練管線其實早已收斂。數(shù)據(jù)加載、損失函數(shù)、優(yōu)化器、評估邏輯、模型導(dǎo)出這些環(huán)節(jié)幾乎是固定的真正變化的只有模型自身以及少數(shù)幾處超參數(shù)。因此完全可以用一份統(tǒng)一代碼模板在 AlexNet、VGG、ResNet、ViT 四個主流架構(gòu)之間橫向切換。這篇文章會圍繞一個可運行的分類項目展開用 PyTorch 實現(xiàn)一個統(tǒng)一模型構(gòu)建層用同一套訓(xùn)練腳本跑通四種架構(gòu)再把訓(xùn)練好的模型導(dǎo)出為 ONNX 做部署驗證。讀完你會擁有一份可以復(fù)用的視覺分類模板后續(xù)換成自己的數(shù)據(jù)和模型結(jié)構(gòu)只需要改少量代碼。1. 這篇文章真正要解決的問題在進(jìn)入代碼之前先把問題說清楚。很多人學(xué)習(xí)視覺模型時傾向于“每個模型單獨學(xué)一遍”但這樣會陷入一個誤區(qū)把大量的時間和注意力放在結(jié)構(gòu)代碼的差異上而忽略了訓(xùn)練流程的共性。AlexNet 和 ViT 雖然看起來差別巨大但它們都可以塞進(jìn)一個標(biāo)準(zhǔn)的監(jiān)督分類流程里讀取圖像、前向傳播、計算交叉熵?fù)p失、反向傳播、更新參數(shù)、周期性評估準(zhǔn)確率。這個模板要解決的三個核心問題是如何不寫四份幾乎一模一樣的訓(xùn)練腳本只維護一份訓(xùn)練代碼就完成多架構(gòu)實驗。如何在四種架構(gòu)之間做橫向?qū)Ρ榷皇敲看味荚诓煌拇a風(fēng)格里來回切換。如何在訓(xùn)練完成后把模型導(dǎo)出為通用格式ONNX讓部署代碼和訓(xùn)練代碼解耦。對讀者來說這篇文章的技術(shù)收益是掌握一個“模型層可插拔、訓(xùn)練層固定、導(dǎo)出層統(tǒng)一”的最小工程范式。如果你是在校學(xué)生可以用它快速完成模型對比實驗如果你是算法工程師可以用它作為新項目的主干腳本如果你剛?cè)腴T深度學(xué)習(xí)可以通過這份模板搞清楚一條完整訓(xùn)練鏈路到底由哪些環(huán)節(jié)組成。這里要先做一個提醒模板的價值在于“統(tǒng)一”但不等于“無腦”。四個架構(gòu)在超參數(shù)敏感度上仍然有明顯差異這一點我會在后面的章節(jié)專門說明。代碼統(tǒng)一解決的是工程重復(fù)問題調(diào)參經(jīng)驗解決的是模型差異問題兩者并不沖突。2. 四個架構(gòu)并不需要四套訓(xùn)練代碼先把四個架構(gòu)放在一起對比這是理解整篇文章的基礎(chǔ)。模型提出年份核心思想典型結(jié)構(gòu)更適用的場景AlexNet2012深度卷積 ReLU Dropout5 層卷積 3 層全連接學(xué)習(xí)基線、新手入門VGG2014小卷積核反復(fù)堆疊多個卷積塊 全連接結(jié)構(gòu)直觀、遷移學(xué)習(xí)基線ResNet2015殘差連接緩解網(wǎng)絡(luò)退化基本塊 / 瓶頸塊通用視覺任務(wù)首選ViT2020圖像切塊 自注意力Patch Embedding Transformer Encoder大數(shù)據(jù) 預(yù)訓(xùn)練微調(diào)AlexNet 的歷史意義在于它證明了“深度”的價值更深的網(wǎng)絡(luò)、非線性的 ReLU 激活、數(shù)據(jù)增強和 Dropout 共同把 ImageNet 分類準(zhǔn)確率提升到了一個此前難以想象的高度。雖然今天直接在生產(chǎn)環(huán)境中使用 AlexNet 的團隊已經(jīng)不多但它仍然是理解卷積神經(jīng)網(wǎng)絡(luò)基本組成的最佳起點。VGG 的核心貢獻(xiàn)是簡化了卷積核的設(shè)計全部使用 3x3 小卷積核通過堆疊卷積塊來增加深度。它把“如何設(shè)計卷積層”變成了“如何堆疊卷基層”這個思路對后來很多網(wǎng)絡(luò)都有影響。缺點也明顯參數(shù)多、計算量大訓(xùn)練和推理成本偏高。ResNet 解決了深度網(wǎng)絡(luò)難以訓(xùn)練的問題。網(wǎng)絡(luò)太深時準(zhǔn)確率反而會下降這叫“退化問題”。ResNet 通過殘差連接讓每一層學(xué)習(xí)的是“殘差”——也就是輸入和輸出之間的差值。這樣梯度可以繞過大部分卷積層直接向后傳播網(wǎng)絡(luò)可以做到很深而不容易訓(xùn)練失敗。在實際工程中ResNet 系列往往是默認(rèn)的視覺基線模型。ViT 則完全改變思路把圖像切成固定大小的 patch映射成 token然后交給 Transformer 的自注意力機制去建模全局關(guān)系。ViT 在超大數(shù)據(jù)集上訓(xùn)練后再微調(diào)效果可以超過 CNN但在數(shù)據(jù)量不足時它不如 CNN 穩(wěn)。理解 ViT 的核心不需要寫完整代碼因為主流框架已經(jīng)提供官方實現(xiàn)你需要做的是把它接入統(tǒng)一訓(xùn)練流程。從工程視角看這四類架構(gòu)的差異集中在“特征提取器”和“分類頭結(jié)構(gòu)”上而訓(xùn)練上下文完全一致。真正應(yīng)該掌握的判斷是學(xué)習(xí)重點要分成兩層——結(jié)構(gòu)層看差異工程層看共性。3. 環(huán)境準(zhǔn)備與版本說明本文代碼基于 PyTorch 生態(tài)示例在 Linux 和 Windows 環(huán)境都可以運行。為了保證代碼不依賴過新或過舊的庫下面列出一個比較穩(wěn)妥的版本區(qū)間。python -c import torch; print(torch.__version__)建議使用Python 3.8 及以上。PyTorch 1.12 及以上推薦 2.x。torchvision 0.13 及以上因為自注意力的 ViT 預(yù)訓(xùn)練模型在 torchvision 中從 0.13 開始提供。onnx 和 onnxruntime用于導(dǎo)出和部署驗證。安裝命令pip install torch torchvision onnx onnxruntime tqdm如果你需要 GPU 加速請根據(jù)自己機器的 CUDA 版本選擇對應(yīng)的 PyTorch 安裝源這一點以 PyTorch 官網(wǎng)給出的安裝命令為準(zhǔn)。CPU 環(huán)境也能運行本文全部代碼只是訓(xùn)練速度會明顯變慢建議先用小 batch size 做驗證。版本問題是最容易踩坑的環(huán)節(jié)。torchvision 在較新版本中把pretrainedTrue改成了weights...的形式本文代碼為了清晰會直接使用新版寫法。如果你使用的是舊版本需要把代碼中的weightsIMAGENET1K_V1換成pretrainedTrue。這個細(xì)節(jié)會在常見問題章節(jié)繼續(xù)展開。4. 數(shù)據(jù)準(zhǔn)備先跑通 CIFAR-10本文示例使用 CIFAR-10 數(shù)據(jù)集。它是視覺分類最常用的教學(xué)數(shù)據(jù)集50000 張訓(xùn)練圖片10000 張測試圖片共 10 個類別。數(shù)據(jù)體積小、下載快特別適合用來驗證模板是否跑通。需要注意一個兩難問題CIFAR-10 原始圖片只有 32x32而 AlexNet、VGG、ResNet、ViT 的標(biāo)準(zhǔn)輸入通常是 224x224。直接把小圖 resize 到大尺寸會損失細(xì)節(jié)也會增加顯存占用但如果完全不 resize某些模型結(jié)構(gòu)可能無法正常工作尤其是依賴固定 patch 數(shù)量的 ViT。模板的處理方式很直接統(tǒng)一 resize 到 224x224。這樣四個模型接受完全相同的輸入代碼路徑是統(tǒng)一的符合本文“模板化”的目標(biāo)。需要理解的是從零訓(xùn)練時把小圖放大到 224 并不理想所以當(dāng)你把模板用到真實數(shù)據(jù)集時應(yīng)該根據(jù)數(shù)據(jù)原始尺寸調(diào)整Resize的數(shù)值。下面是數(shù)據(jù)加載部分的代碼。這里只寫一份訓(xùn)練集和驗證集都使用同一套邏輯區(qū)別在于是否做隨機數(shù)據(jù)增強。# 文件路徑data_utils.py import torchvision.transforms as transforms def get_transform(train: bool True): normalize transforms.Normalize( mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225] ) if train: return transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.2, contrast0.2), transforms.ToTensor(), normalize, ]) return transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), normalize, ])訓(xùn)練階段加入隨機水平翻轉(zhuǎn)和顏色抖動是為了提高模型泛化能力。驗證階段不做隨機增強只做 resize 和歸一化保證評估結(jié)果的穩(wěn)定性。在train.py中加載數(shù)據(jù)時使用torchvision.datasets.CIFAR10并設(shè)置downloadTrue。第一次運行會自動下載數(shù)據(jù)之后會直接讀取本地緩存。數(shù)據(jù)規(guī)模較小所以 DataLoader 的num_workers設(shè)為 2 即可過大會在 Windows 上引發(fā)額外問題。5. 統(tǒng)一模型構(gòu)建層用注冊表管理四種架構(gòu)很多人會問為什么不用簡單的 if-else 來切換模型if-else 當(dāng)然能跑但隨著模型種類增多代碼會越來越混亂。更好的做法是用一個小的注冊表模式每個模型由一個 builder 函數(shù)負(fù)責(zé)構(gòu)建函數(shù)名和模型字符串綁定。以后新增模型只需要加一個函數(shù)不需要改調(diào)用方。下面這段代碼是模板的關(guān)鍵部分。它封裝了四個模型的構(gòu)建邏輯并且支持從 ImageNet 預(yù)訓(xùn)練權(quán)重初始化。# 文件路徑models.py import torch.nn as nn import torchvision _MODEL_FACTORY {} def register_model(name: str): def decorator(fn): _MODEL_FACTORY[name] fn return fn return decorator def build_model(name: str, num_classes: int 10, pretrained: bool False): if name not in _MODEL_FACTORY: raise KeyError(f未知模型{name}可選模型 {list(_MODEL_FACTORY.keys())}) return _MODEL_FACTORY[name](num_classesnum_classes, pretrainedpretrained) register_model(alexnet) def build_alexnet(num_classes: int, pretrained: bool False): weights IMAGENET1K_V1 if pretrained else None model torchvision.models.alexnet(weightsweights) model.classifier[6] nn.Linear(model.classifier[6].in_features, num_classes) return model register_model(vgg13) def build_vgg13(num_classes: int, pretrained: bool False): weights IMAGENET1K_V1 if pretrained else None model torchvision.models.vgg13(weightsweights) model.classifier[6] nn.Linear(model.classifier[6].in_features, num_classes) return model register_model(resnet18) def build_resnet18(num_classes: int, pretrained: bool False): weights IMAGENET1K_V1 if pretrained else None model torchvision.models.resnet18(weightsweights) model.fc nn.Linear(model.fc.in_features, num_classes) return model register_model(vit_b_16) def build_vit_b_16(num_classes: int, pretrained: bool False): weights IMAGENET1K_V1 if pretrained else None model torchvision.models.vit_b_16(weightsweights) model.heads.head nn.Linear(model.heads.head.in_features, num_classes) return model這段代碼里最重要的一層含義是模型構(gòu)建是“可插拔”的。調(diào)用方永遠(yuǎn)只依賴build_model(name, num_classes, pretrained)至于 name 對應(yīng)的是卷積網(wǎng)絡(luò)還是 Transformer調(diào)用方不需要關(guān)心。這樣訓(xùn)練腳本可以徹底固定下來不會因為換了模型而改動主流程。如果你使用的是較老版本的 torchvisionweights參數(shù)不存在。更保守的做法是把weights改成pretrained但那樣會失去新版本 API 的優(yōu)勢。這里建議優(yōu)先升級 torchvision而不是為了兼容舊版本把代碼寫得更復(fù)雜。6. 訓(xùn)練循環(huán)一份腳本跑通四個模型訓(xùn)練主腳本是整個模板的核心。它做的事情是所有分類任務(wù)共用的加載數(shù)據(jù)、構(gòu)建模型、定義損失函數(shù)和優(yōu)化器、迭代若干輪、每輪驗證并保存最佳權(quán)重。下面給出完整訓(xùn)練代碼。這是一個最小可運行版本沒有引入分布式訓(xùn)練和復(fù)雜的日志框架方便理解主干邏輯。# 文件路徑train.py import argparse import os import torch import torch.nn as nn from torch.utils.data import DataLoader from torchvision.datasets import CIFAR10 from data_utils import get_transform from models import build_model def train_one_epoch(model, dataloader, criterion, optimizer, device): model.train() total_loss, total_correct, total_num 0.0, 0, 0 for images, labels in dataloader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() total_loss loss.item() * images.size(0) total_correct (outputs.argmax(dim1) labels).sum().item() total_num images.size(0) return total_loss / total_num, total_correct / total_num torch.no_grad() def evaluate(model, dataloader, criterion, device): model.eval() total_loss, total_correct, total_num 0.0, 0, 0 for images, labels in dataloader: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) total_loss loss.item() * images.size(0) total_correct (outputs.argmax(dim1) labels).sum().item() total_num images.size(0) return total_loss / total_num, total_correct / total_num def main(): parser argparse.ArgumentParser() parser.add_argument(--model, typestr, defaultresnet18, choices[alexnet, vgg13, resnet18, vit_b_16]) parser.add_argument(--epochs, typeint, default30) parser.add_argument(--batch-size, typeint, default64) parser.add_argument(--lr, typefloat, default1e-2) parser.add_argument(--optimizer, typestr, defaultsgd, choices[sgd, adamw]) parser.add_argument(--pretrained, actionstore_true) parser.add_argument(--device, typestr, defaultcuda if torch.cuda.is_available() else cpu) parser.add_argument(--data-dir, typestr, default./data) parser.add_argument(--save-dir, typestr, default./checkpoints) args parser.parse_args() os.makedirs(args.save_dir, exist_okTrue) device torch.device(args.device) train_loader DataLoader( CIFAR10(rootargs.data_dir, trainTrue, transformget_transform(True), downloadTrue), batch_sizeargs.batch_size, shuffleTrue, num_workers2, pin_memoryTrue ) val_loader DataLoader( CIFAR10(rootargs.data_dir, trainFalse, transformget_transform(False), downloadTrue), batch_sizeargs.batch_size, shuffleFalse, num_workers2, pin_memoryTrue ) model build_model(args.model, num_classes10, pretrainedargs.pretrained).to(device) criterion nn.CrossEntropyLoss() if args.optimizer sgd: optimizer torch.optim.SGD(model.parameters(), lrargs.lr, momentum0.9, weight_decay5e-4) else: optimizer torch.optim.AdamW(model.parameters(), lrargs.lr, weight_decay5e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_maxargs.epochs) best_acc 0.0 for epoch in range(args.epochs): train_loss, train_acc train_one_epoch(model, train_loader, criterion, optimizer, device) val_loss, val_acc evaluate(model, val_loader, criterion, device) scheduler.step() print(fEpoch {epoch 1}/{args.epochs} | ftrain loss {train_loss:.4f} acc {train_acc * 100:.1f}% | fval acc {val_acc * 100:.1f}% | flr {scheduler.get_last_lr()[0]:.2e}) if val_acc best_acc: best_acc val_acc torch.save({ model: model.state_dict(), epoch: epoch, val_acc: best_acc, }, os.path.join(args.save_dir, f{args.model}_best.pt)) print(f訓(xùn)練完成最佳驗證準(zhǔn)確率{best_acc * 100:.2f}%) print(f權(quán)重已保存至{os.path.join(args.save_dir, f{args.model}_best.pt)}) if __name__ __main__: main()這段代碼的邏輯并不復(fù)雜但有三個地方需要重點看優(yōu)化器支持 SGD 和 AdamW。ResNet、VGG 使用帶動量的 SGD 通常更穩(wěn)ViT 更推薦 AdamW配合較小學(xué)習(xí)率。學(xué)習(xí)率調(diào)度使用余弦退火。它能在訓(xùn)練后期逐步降低學(xué)習(xí)率提升收斂穩(wěn)定性。每輪結(jié)束后立即在驗證集上評估并保存驗證準(zhǔn)確率最高的權(quán)重。這樣訓(xùn)練中斷時不至于丟失最佳結(jié)果。如果你想繼續(xù)跑第二個模型不需要改動任何代碼只需換--model參數(shù)。訓(xùn)練腳本本身不感知模型內(nèi)部結(jié)構(gòu)差異這正是模板化的價值所在。7. 運行與驗證從命令行到日志判斷在開始訓(xùn)練之前先用最小的配置驗證模板本身是否工作正常。這里建議先把 batch size 調(diào)小比如 16確定數(shù)據(jù)加載和模型 forward 沒有問題再跑正式訓(xùn)練。四個模型分別通過以下命令啟動# AlexNet從零訓(xùn)練 python train.py --model alexnet --epochs 30 --batch-size 64 --lr 1e-2 # VGG13從零訓(xùn)練 python train.py --model vgg13 --epochs 30 --batch-size 64 --lr 1e-2 # ResNet18從零訓(xùn)練 python train.py --model resnet18 --epochs 30 --batch-size 64 --lr 1e-2 # ViT使用 AdamW 并且加載 ImageNet 預(yù)訓(xùn)練權(quán)重推薦 python train.py --model vit_b_16 --epochs 30 --batch-size 64 --lr 1e-4 --optimizer adamw --pretrained如果你機器顯存有限可以繼續(xù)降低 batch size。Batch size 不影響代碼結(jié)構(gòu)只影響訓(xùn)練速度和穩(wěn)定性。正常運行時的輸出格式類似這樣Epoch 1/30 | train loss 1.8935 acc 31.2% | val acc 47.1% | lr 1.00e-02 Epoch 2/30 | train loss 1.4102 acc 47.5% | val acc 55.8% | lr 9.66e-03如何判斷訓(xùn)練是否健康先看 train loss 是否持續(xù)下降。如果 train loss 下降但 val acc 不漲說明過擬合如果 train loss 和 val acc 都不動大概率是學(xué)習(xí)率設(shè)置有問題或者數(shù)據(jù)加載異常。不同模型在這個模板上的表現(xiàn)會遵循一些規(guī)律。ResNet18 從零訓(xùn)練到 30 輪驗證準(zhǔn)確率通常能接近 80% 甚至更高比較穩(wěn)定。AlexNet 參數(shù)多但結(jié)構(gòu)相對淺在 CIFAR-10 這種小圖上提升較慢。VGG13 訓(xùn)練速度慢顯存占用高但準(zhǔn)確率通常不錯。ViT 如果不做預(yù)訓(xùn)練在 CIFAR-10 上很容易過擬合這也是為什么它在小數(shù)據(jù)集上口碑不如 CNN 的直接原因。穿針引線地說模板讓你快速得到結(jié)果但理解不同模型的行為差異仍然要靠實驗對比和調(diào)參經(jīng)驗。8. 導(dǎo)出 ONNX訓(xùn)練與部署的解耦模型訓(xùn)練完成后下一步通常是部署。很多人直接把 PyTorch 模型放到服務(wù)里這樣雖然能用但推理代碼和訓(xùn)練代碼耦合很重。更通用的做法是導(dǎo)出為 ONNX再用 ONNX Runtime 這種專門的推理引擎加載。ONNX 的好處是與框架無關(guān)后續(xù)在 CPU、GPU 甚至邊緣設(shè)備上都能找到對應(yīng)的運行環(huán)境。下面這段代碼把訓(xùn)練好的 checkpoint 導(dǎo)出為 ONNX 文件。# 文件路徑export_onnx.py import argparse import torch from models import build_model def main(): parser argparse.ArgumentParser() parser.add_argument(--model, typestr, defaultresnet18, choices[alexnet, vgg13, resnet18, vit_b_16]) parser.add_argument(--weight, typestr, requiredTrue) parser.add_argument(--output, typestr, defaultmodel.onnx) parser.add_argument(--size, typeint, default224) args parser.parse_args() model build_model(args.model, num_classes10, pretrainedFalse) state_dict torch.load(args.weight, map_locationcpu) model.load_state_dict(state_dict[model]) model.eval() dummy_input torch.randn(1, 3, args.size, args.size) torch.onnx.export( model, dummy_input, args.output, input_names[input], output_names[logits], dynamic_axes{input: {0: batch}, logits: {0: batch}}, opset_version17, ) print(fONNX 導(dǎo)出完成{args.output}) if __name__ __main__: main()導(dǎo)出命令python export_onnx.py --model resnet18 --weight checkpoints/resnet18_best.pt --output resnet18.onnx導(dǎo)出 ONNX 時最關(guān)鍵的是dynamic_axes。如果不設(shè)置ONNX 模型輸入 batch 會被固定為 1部署時只能一次推理一張圖。設(shè)置dynamic_axes后batch 維度可以是動態(tài)的更加靈活。導(dǎo)出完成后用 ONNX Runtime 做一次推理驗證確認(rèn)導(dǎo)出沒有損壞模型結(jié)構(gòu)。# 文件路徑onnx_infer.py import numpy as np import onnxruntime as ort sess ort.InferenceSession(resnet18.onnx) input_name sess.get_inputs()[0].name x np.random.randn(1, 3, 224, 224).astype(np.float32) logits sess.run(None, {input_name: x})[0] print(輸出 shape:, logits.shape) print(預(yù)測類別索引:, logits.argmax(axis1))這一步的輸出應(yīng)該是一個長度為 10 的向量因為 CIFAR-10 有 10 個類別。如果你導(dǎo)出的模型分類頭是 1000 類而不是 10 類說明導(dǎo)出腳本構(gòu)建模型時沒有替換分類頭或者加載的權(quán)重和模型結(jié)構(gòu)不匹配。這里也能看出模板設(shè)計的另一個收益訓(xùn)練權(quán)重文件里保存的是“模型參數(shù)”而不是完整模型對象導(dǎo)出腳本只依賴build_model這個名字來重建結(jié)構(gòu)。所以模型定義層必須保持穩(wěn)定這也是前面使用注冊表模式的原因。9. 常見問題與排查思路在實際運行中讀者最常遇到的并不是模型理論問題而是環(huán)境、版本、顯存這類工程問題。下面整理了一份排查表按優(yōu)先級排列。問題現(xiàn)象可能原因排查方式解決方案導(dǎo)入vit_b_16報 ImportErrortorchvision 版本過低沒有 ViT 實現(xiàn)執(zhí)行python -c import torchvision; print(torchvision.__version__)升級 torchvision 到 0.13 以上weights參數(shù)不被識別代碼使用了新版 API舊版只支持pretrainedTrue查看安裝的 PyTorch 版本對應(yīng)文檔改為pretrainedTrue或升級版本顯存溢出OOMbatch size 過大或 224x224 輸入在 ViT 上顯存消耗高觀察日志中哪個階段報錯降低 batch size 重試--batch-size 16或更小并減少num_workerstrain loss 不下降學(xué)習(xí)率過大或過小也可能是數(shù)據(jù)沒有正確歸一化打日志檢查 loss 數(shù)量級嘗試 lr 1e-3從零訓(xùn)練用 SGD lr1e-2預(yù)訓(xùn)練微調(diào)用 lr1e-3ViT 在 CIFAR-10 上效果差小數(shù)據(jù)量 無預(yù)訓(xùn)練Transformer 容易過擬合觀察 val acc 是否遠(yuǎn)低于 train acc使用--pretrained或換用 ResNet 作為基線ONNX 導(dǎo)出后 batch 固定為 1dynamic_axes沒有正確設(shè)置用 Netron 打開 ONNX 文件查看輸入 shape在torch.onnx.export中設(shè)置dynamic_axes加載權(quán)重時 shape 不匹配checkpoint 的模型和當(dāng)前--model不一致對比訓(xùn)練命令和導(dǎo)出命令中的模型名確保訓(xùn)練和導(dǎo)出使用同一個模型名稱驗證準(zhǔn)確率高于訓(xùn)練準(zhǔn)確率Dropout 和 BatchNorm 在 train/eval 模式行為不同檢查model.train()和model.eval()是否正確切換在推理和驗證前必須執(zhí)行model.eval()這張表里的多數(shù)問題都無法靠“重裝環(huán)境”解決而是源于對版本差異和模型行為的不熟悉。建議你在開始訓(xùn)練前先確認(rèn)三件事torchvision 版本足夠新數(shù)據(jù)集能夠正常下載CIFAR-10 的歸一化參數(shù)已經(jīng)配置正確。10. 工程實踐建議與超參選擇經(jīng)驗?zāi)0迥茏屇憧焖侔阉心P团芷饋淼谡鎸嶍椖恐惺褂脮r還需要遵守一些工程經(jīng)驗否則很容易被表面上的“跑通”欺騙。第一先小后大。任何新的數(shù)據(jù)集接入模板都不要直接全量訓(xùn)練。先用每個類別 100 到 200 張圖的子集跑 1 到 2 個 epoch確認(rèn)數(shù)據(jù)加載正確、loss 能下降、模型可以收斂到隨機水平以上。工程里大量 bug 都是在全量訓(xùn)練跑了幾小時后才暴露的小規(guī)模驗證成本低得多。第二學(xué)習(xí)率策略要區(qū)分模型。ResNet 和 VGG 使用 SGD momentum 是經(jīng)過大量實踐驗證的穩(wěn)妥選擇學(xué)習(xí)率從 1e-2 開始調(diào)。ViT 更推薦 AdamW配合 warmup 和較小的學(xué)習(xí)率比如 1e-4 到 5e-4。如果套用同一個學(xué)習(xí)率到所有模型不是模板的錯而是沒有理解模型差異的體現(xiàn)。第三數(shù)據(jù)處理比模型結(jié)構(gòu)更影響最終結(jié)果。CIFAR-10 因為圖小簡單增強就夠用但真實數(shù)據(jù)往往需要做 resize、隨機裁剪、旋轉(zhuǎn)、亮度擾動等多種組合。當(dāng)模板遷移到新數(shù)據(jù)時首先要調(diào)整的是get_transform而不是急著換更復(fù)雜的模型。第四checkpoint 里不僅保存模型權(quán)重還可以保存優(yōu)化器狀態(tài)、當(dāng)前 epoch、學(xué)習(xí)率等信息。本文示例為了簡潔只保存了模型權(quán)重和驗證準(zhǔn)確率工程中建議至少加上優(yōu)化器狀態(tài)便于中斷后從最近的狀態(tài)繼續(xù)訓(xùn)練。第五從零訓(xùn)練還是微調(diào)預(yù)訓(xùn)練權(quán)重必須想清楚。大數(shù)據(jù)量任務(wù)可以從零訓(xùn)練但視覺任務(wù)普遍推薦使用 ImageNet 預(yù)訓(xùn)練權(quán)重做初始化尤其是 ViT。預(yù)訓(xùn)練初始化相當(dāng)于讓模型從“已經(jīng)會看邊的表示”開始而不是從“純隨機”開始這在中小數(shù)據(jù)集上差別非常大。第六部署前要檢查eval()狀態(tài)。PyTorch 模型在導(dǎo)出 ONNX 前一般要求調(diào)用model.eval()因為 BatchNorm 和 Dropout 在訓(xùn)練與推理時的行為不同。如果不做導(dǎo)出的模型行為可能不穩(wěn)定。11. 總結(jié)與后續(xù)學(xué)習(xí)方向這份模板的核心價值不在于代碼技巧而在于把視覺分類任務(wù)的工程鏈路固化下來數(shù)據(jù)、模型、訓(xùn)練、評估、導(dǎo)出。用同一份代碼跑 AlexNet、VGG、ResNet、ViT本質(zhì)上是承認(rèn)一個事實——這些模型的差異在結(jié)構(gòu)設(shè)計層面而工程編排層面已經(jīng)標(biāo)準(zhǔn)化。理解了這一點你就可以把精力放到真正值得研究的地方。下一步建議從這個方向開始實踐先把模板在 CIFAR-10 上完整跑一遍重點觀察 ResNet 和 ViT 在相同數(shù)據(jù)下的行為差異然后用自己的數(shù)據(jù)集替換 CIFAR-10按實際分類數(shù)量修改build_model中的num_classes。如果數(shù)據(jù)量較大可以嘗試使用混合精度訓(xùn)練如果模型較大可以考慮分布式訓(xùn)練如果涉及目標(biāo)檢測或語義分割則需要在本文模板基礎(chǔ)上增加對應(yīng)任務(wù)頭。最后提醒一句模板是有適用邊界的。分類任務(wù)的訓(xùn)練管線收斂了不代表所有視覺任務(wù)都適合這個套路。當(dāng)你進(jìn)入檢測、分割、多模態(tài)等方向時可以復(fù)用本文“注冊表管理模型、固定訓(xùn)練循環(huán)、統(tǒng)一導(dǎo)出格式”的思維方式但代碼結(jié)構(gòu)需要重新設(shè)計。把這份模板收藏起來作為你快速驗證模型和建立實驗基線的起點它會幫你省下大量重復(fù)造輪子的時間。