很多剛接觸視覺算法的同學都會遇到一個類似困境在文章或課程里把 AlexNet、VGG、ResNet、ViT 的結構圖看得明明白白真正打開編輯器準備訓練時卻發(fā)現自己只有“網絡結構代碼”沒有一套完整的數據加載、訓練循環(huán)、評估和導出腳本。于是從 A 倉庫復制數據增強從 B 倉庫找訓練參數再從 C 倉庫拼接模型調用最后訓練到一半才發(fā)現學習率策略和優(yōu)化器根本不匹配。這不是個人能力問題而是當前視覺入門教程普遍存在的空白結構講解很多工程閉環(huán)太少?？梢越o出一個明確判斷在圖像分類這個方向上從經典 CNN 到 Transformer訓練管線其實早已收斂。數據加載、損失函數、優(yōu)化器、評估邏輯、模型導出這些環(huán)節(jié)幾乎是固定的真正變化的只有模型自身以及少數幾處超參數。因此完全可以用一份統一代碼模板在 AlexNet、VGG、ResNet、ViT 四個主流架構之間橫向切換。這篇文章會圍繞一個可運行的分類項目展開用 PyTorch 實現一個統一模型構建層用同一套訓練腳本跑通四種架構再把訓練好的模型導出為 ONNX 做部署驗證。讀完你會擁有一份可以復用的視覺分類模板后續(xù)換成自己的數據和模型結構只需要改少量代碼。1. 這篇文章真正要解決的問題在進入代碼之前先把問題說清楚。很多人學習視覺模型時傾向于“每個模型單獨學一遍”但這樣會陷入一個誤區(qū)把大量的時間和注意力放在結構代碼的差異上而忽略了訓練流程的共性。AlexNet 和 ViT 雖然看起來差別巨大但它們都可以塞進一個標準的監(jiān)督分類流程里讀取圖像、前向傳播、計算交叉熵損失、反向傳播、更新參數、周期性評估準確率。這個模板要解決的三個核心問題是如何不寫四份幾乎一模一樣的訓練腳本只維護一份訓練代碼就完成多架構實驗。如何在四種架構之間做橫向對比而不是每次都在不同的代碼風格里來回切換。如何在訓練完成后把模型導出為通用格式ONNX讓部署代碼和訓練代碼解耦。對讀者來說這篇文章的技術收益是掌握一個“模型層可插拔、訓練層固定、導出層統一”的最小工程范式。如果你是在校學生可以用它快速完成模型對比實驗如果你是算法工程師可以用它作為新項目的主干腳本如果你剛入門深度學習可以通過這份模板搞清楚一條完整訓練鏈路到底由哪些環(huán)節(jié)組成。這里要先做一個提醒模板的價值在于“統一”但不等于“無腦”。四個架構在超參數敏感度上仍然有明顯差異這一點我會在后面的章節(jié)專門說明。代碼統一解決的是工程重復問題調參經驗解決的是模型差異問題兩者并不沖突。2. 四個架構并不需要四套訓練代碼先把四個架構放在一起對比這是理解整篇文章的基礎。模型提出年份核心思想典型結構更適用的場景AlexNet2012深度卷積 ReLU Dropout5 層卷積 3 層全連接學習基線、新手入門VGG2014小卷積核反復堆疊多個卷積塊 全連接結構直觀、遷移學習基線ResNet2015殘差連接緩解網絡退化基本塊 / 瓶頸塊通用視覺任務首選ViT2020圖像切塊 自注意力Patch Embedding Transformer Encoder大數據 預訓練微調AlexNet 的歷史意義在于它證明了“深度”的價值更深的網絡、非線性的 ReLU 激活、數據增強和 Dropout 共同把 ImageNet 分類準確率提升到了一個此前難以想象的高度。雖然今天直接在生產環(huán)境中使用 AlexNet 的團隊已經不多但它仍然是理解卷積神經網絡基本組成的最佳起點。VGG 的核心貢獻是簡化了卷積核的設計全部使用 3x3 小卷積核通過堆疊卷積塊來增加深度。它把“如何設計卷積層”變成了“如何堆疊卷基層”這個思路對后來很多網絡都有影響。缺點也明顯參數多、計算量大訓練和推理成本偏高。ResNet 解決了深度網絡難以訓練的問題。網絡太深時準確率反而會下降這叫“退化問題”。ResNet 通過殘差連接讓每一層學習的是“殘差”——也就是輸入和輸出之間的差值。這樣梯度可以繞過大部分卷積層直接向后傳播網絡可以做到很深而不容易訓練失敗。在實際工程中ResNet 系列往往是默認的視覺基線模型。ViT 則完全改變思路把圖像切成固定大小的 patch映射成 token然后交給 Transformer 的自注意力機制去建模全局關系。ViT 在超大數據集上訓練后再微調效果可以超過 CNN但在數據量不足時它不如 CNN 穩(wěn)。理解 ViT 的核心不需要寫完整代碼因為主流框架已經提供官方實現你需要做的是把它接入統一訓練流程。從工程視角看這四類架構的差異集中在“特征提取器”和“分類頭結構”上而訓練上下文完全一致。真正應該掌握的判斷是學習重點要分成兩層——結構層看差異工程層看共性。3. 環(huá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 預訓練模型在 torchvision 中從 0.13 開始提供。onnx 和 onnxruntime用于導出和部署驗證。安裝命令pip install torch torchvision onnx onnxruntime tqdm如果你需要 GPU 加速請根據自己機器的 CUDA 版本選擇對應的 PyTorch 安裝源這一點以 PyTorch 官網給出的安裝命令為準。CPU 環(huán)境也能運行本文全部代碼只是訓練速度會明顯變慢建議先用小 batch size 做驗證。版本問題是最容易踩坑的環(huán)節(jié)。torchvision 在較新版本中把pretrainedTrue改成了weights...的形式本文代碼為了清晰會直接使用新版寫法。如果你使用的是舊版本需要把代碼中的weightsIMAGENET1K_V1換成pretrainedTrue。這個細節(jié)會在常見問題章節(jié)繼續(xù)展開。4. 數據準備先跑通 CIFAR-10本文示例使用 CIFAR-10 數據集。它是視覺分類最常用的教學數據集50000 張訓練圖片10000 張測試圖片共 10 個類別。數據體積小、下載快特別適合用來驗證模板是否跑通。需要注意一個兩難問題CIFAR-10 原始圖片只有 32x32而 AlexNet、VGG、ResNet、ViT 的標準輸入通常是 224x224。直接把小圖 resize 到大尺寸會損失細節(jié)也會增加顯存占用但如果完全不 resize某些模型結構可能無法正常工作尤其是依賴固定 patch 數量的 ViT。模板的處理方式很直接統一 resize 到 224x224。這樣四個模型接受完全相同的輸入代碼路徑是統一的符合本文“模板化”的目標。需要理解的是從零訓練時把小圖放大到 224 并不理想所以當你把模板用到真實數據集時應該根據數據原始尺寸調整Resize的數值。下面是數據加載部分的代碼。這里只寫一份訓練集和驗證集都使用同一套邏輯區(qū)別在于是否做隨機數據增強。# 文件路徑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, ])訓練階段加入隨機水平翻轉和顏色抖動是為了提高模型泛化能力。驗證階段不做隨機增強只做 resize 和歸一化保證評估結果的穩(wěn)定性。在train.py中加載數據時使用torchvision.datasets.CIFAR10并設置downloadTrue。第一次運行會自動下載數據之后會直接讀取本地緩存。數據規(guī)模較小所以 DataLoader 的num_workers設為 2 即可過大會在 Windows 上引發(fā)額外問題。5. 統一模型構建層用注冊表管理四種架構很多人會問為什么不用簡單的 if-else 來切換模型if-else 當然能跑但隨著模型種類增多代碼會越來越混亂。更好的做法是用一個小的注冊表模式每個模型由一個 builder 函數負責構建函數名和模型字符串綁定。以后新增模型只需要加一個函數不需要改調用方。下面這段代碼是模板的關鍵部分。它封裝了四個模型的構建邏輯并且支持從 ImageNet 預訓練權重初始化。# 文件路徑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這段代碼里最重要的一層含義是模型構建是“可插拔”的。調用方永遠只依賴build_model(name, num_classes, pretrained)至于 name 對應的是卷積網絡還是 Transformer調用方不需要關心。這樣訓練腳本可以徹底固定下來不會因為換了模型而改動主流程。如果你使用的是較老版本的 torchvisionweights參數不存在。更保守的做法是把weights改成pretrained但那樣會失去新版本 API 的優(yōu)勢。這里建議優(yōu)先升級 torchvision而不是為了兼容舊版本把代碼寫得更復雜。6. 訓練循環(huán)一份腳本跑通四個模型訓練主腳本是整個模板的核心。它做的事情是所有分類任務共用的加載數據、構建模型、定義損失函數和優(yōu)化器、迭代若干輪、每輪驗證并保存最佳權重。下面給出完整訓練代碼。這是一個最小可運行版本沒有引入分布式訓練和復雜的日志框架方便理解主干邏輯。# 文件路徑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訓練完成最佳驗證準確率{best_acc * 100:.2f}%) print(f權重已保存至{os.path.join(args.save_dir, f{args.model}_best.pt)}) if __name__ __main__: main()這段代碼的邏輯并不復雜但有三個地方需要重點看優(yōu)化器支持 SGD 和 AdamW。ResNet、VGG 使用帶動量的 SGD 通常更穩(wěn)ViT 更推薦 AdamW配合較小學習率。學習率調度使用余弦退火。它能在訓練后期逐步降低學習率提升收斂穩(wěn)定性。每輪結束后立即在驗證集上評估并保存驗證準確率最高的權重。這樣訓練中斷時不至于丟失最佳結果。如果你想繼續(xù)跑第二個模型不需要改動任何代碼只需換--model參數。訓練腳本本身不感知模型內部結構差異這正是模板化的價值所在。7. 運行與驗證從命令行到日志判斷在開始訓練之前先用最小的配置驗證模板本身是否工作正常。這里建議先把 batch size 調小比如 16確定數據加載和模型 forward 沒有問題再跑正式訓練。四個模型分別通過以下命令啟動# AlexNet從零訓練 python train.py --model alexnet --epochs 30 --batch-size 64 --lr 1e-2 # VGG13從零訓練 python train.py --model vgg13 --epochs 30 --batch-size 64 --lr 1e-2 # ResNet18從零訓練 python train.py --model resnet18 --epochs 30 --batch-size 64 --lr 1e-2 # ViT使用 AdamW 并且加載 ImageNet 預訓練權重推薦 python train.py --model vit_b_16 --epochs 30 --batch-size 64 --lr 1e-4 --optimizer adamw --pretrained如果你機器顯存有限可以繼續(xù)降低 batch size。Batch size 不影響代碼結構只影響訓練速度和穩(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如何判斷訓練是否健康先看 train loss 是否持續(xù)下降。如果 train loss 下降但 val acc 不漲說明過擬合如果 train loss 和 val acc 都不動大概率是學習率設置有問題或者數據加載異常。不同模型在這個模板上的表現會遵循一些規(guī)律。ResNet18 從零訓練到 30 輪驗證準確率通常能接近 80% 甚至更高比較穩(wěn)定。AlexNet 參數多但結構相對淺在 CIFAR-10 這種小圖上提升較慢。VGG13 訓練速度慢顯存占用高但準確率通常不錯。ViT 如果不做預訓練在 CIFAR-10 上很容易過擬合這也是為什么它在小數據集上口碑不如 CNN 的直接原因。穿針引線地說模板讓你快速得到結果但理解不同模型的行為差異仍然要靠實驗對比和調參經驗。8. 導出 ONNX訓練與部署的解耦模型訓練完成后下一步通常是部署。很多人直接把 PyTorch 模型放到服務里這樣雖然能用但推理代碼和訓練代碼耦合很重。更通用的做法是導出為 ONNX再用 ONNX Runtime 這種專門的推理引擎加載。ONNX 的好處是與框架無關后續(xù)在 CPU、GPU 甚至邊緣設備上都能找到對應的運行環(huán)境。下面這段代碼把訓練好的 checkpoint 導出為 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 導出完成{args.output}) if __name__ __main__: main()導出命令python export_onnx.py --model resnet18 --weight checkpoints/resnet18_best.pt --output resnet18.onnx導出 ONNX 時最關鍵的是dynamic_axes。如果不設置ONNX 模型輸入 batch 會被固定為 1部署時只能一次推理一張圖。設置dynamic_axes后batch 維度可以是動態(tài)的更加靈活。導出完成后用 ONNX Runtime 做一次推理驗證確認導出沒有損壞模型結構。# 文件路徑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(預測類別索引:, logits.argmax(axis1))這一步的輸出應該是一個長度為 10 的向量因為 CIFAR-10 有 10 個類別。如果你導出的模型分類頭是 1000 類而不是 10 類說明導出腳本構建模型時沒有替換分類頭或者加載的權重和模型結構不匹配。這里也能看出模板設計的另一個收益訓練權重文件里保存的是“模型參數”而不是完整模型對象導出腳本只依賴build_model這個名字來重建結構。所以模型定義層必須保持穩(wěn)定這也是前面使用注冊表模式的原因。9. 常見問題與排查思路在實際運行中讀者最常遇到的并不是模型理論問題而是環(huán)境、版本、顯存這類工程問題。下面整理了一份排查表按優(yōu)先級排列。問題現象可能原因排查方式解決方案導入vit_b_16報 ImportErrortorchvision 版本過低沒有 ViT 實現執(zhí)行python -c import torchvision; print(torchvision.__version__)升級 torchvision 到 0.13 以上weights參數不被識別代碼使用了新版 API舊版只支持pretrainedTrue查看安裝的 PyTorch 版本對應文檔改為pretrainedTrue或升級版本顯存溢出OOMbatch size 過大或 224x224 輸入在 ViT 上顯存消耗高觀察日志中哪個階段報錯降低 batch size 重試--batch-size 16或更小并減少num_workerstrain loss 不下降學習率過大或過小也可能是數據沒有正確歸一化打日志檢查 loss 數量級嘗試 lr 1e-3從零訓練用 SGD lr1e-2預訓練微調用 lr1e-3ViT 在 CIFAR-10 上效果差小數據量 無預訓練Transformer 容易過擬合觀察 val acc 是否遠低于 train acc使用--pretrained或換用 ResNet 作為基線ONNX 導出后 batch 固定為 1dynamic_axes沒有正確設置用 Netron 打開 ONNX 文件查看輸入 shape在torch.onnx.export中設置dynamic_axes加載權重時 shape 不匹配checkpoint 的模型和當前--model不一致對比訓練命令和導出命令中的模型名確保訓練和導出使用同一個模型名稱驗證準確率高于訓練準確率Dropout 和 BatchNorm 在 train/eval 模式行為不同檢查model.train()和model.eval()是否正確切換在推理和驗證前必須執(zhí)行model.eval()這張表里的多數問題都無法靠“重裝環(huán)境”解決而是源于對版本差異和模型行為的不熟悉。建議你在開始訓練前先確認三件事torchvision 版本足夠新數據集能夠正常下載CIFAR-10 的歸一化參數已經配置正確。10. 工程實踐建議與超參選擇經驗模板能讓你快速把所有模型跑起來但在真實項目中使用時還需要遵守一些工程經驗否則很容易被表面上的“跑通”欺騙。第一先小后大。任何新的數據集接入模板都不要直接全量訓練。先用每個類別 100 到 200 張圖的子集跑 1 到 2 個 epoch確認數據加載正確、loss 能下降、模型可以收斂到隨機水平以上。工程里大量 bug 都是在全量訓練跑了幾小時后才暴露的小規(guī)模驗證成本低得多。第二學習率策略要區(qū)分模型。ResNet 和 VGG 使用 SGD momentum 是經過大量實踐驗證的穩(wěn)妥選擇學習率從 1e-2 開始調。ViT 更推薦 AdamW配合 warmup 和較小的學習率比如 1e-4 到 5e-4。如果套用同一個學習率到所有模型不是模板的錯而是沒有理解模型差異的體現。第三數據處理比模型結構更影響最終結果。CIFAR-10 因為圖小簡單增強就夠用但真實數據往往需要做 resize、隨機裁剪、旋轉、亮度擾動等多種組合。當模板遷移到新數據時首先要調整的是get_transform而不是急著換更復雜的模型。第四checkpoint 里不僅保存模型權重還可以保存優(yōu)化器狀態(tài)、當前 epoch、學習率等信息。本文示例為了簡潔只保存了模型權重和驗證準確率工程中建議至少加上優(yōu)化器狀態(tài)便于中斷后從最近的狀態(tài)繼續(xù)訓練。第五從零訓練還是微調預訓練權重必須想清楚。大數據量任務可以從零訓練但視覺任務普遍推薦使用 ImageNet 預訓練權重做初始化尤其是 ViT。預訓練初始化相當于讓模型從“已經會看邊的表示”開始而不是從“純隨機”開始這在中小數據集上差別非常大。第六部署前要檢查eval()狀態(tài)。PyTorch 模型在導出 ONNX 前一般要求調用model.eval()因為 BatchNorm 和 Dropout 在訓練與推理時的行為不同。如果不做導出的模型行為可能不穩(wěn)定。11. 總結與后續(xù)學習方向這份模板的核心價值不在于代碼技巧而在于把視覺分類任務的工程鏈路固化下來數據、模型、訓練、評估、導出。用同一份代碼跑 AlexNet、VGG、ResNet、ViT本質上是承認一個事實——這些模型的差異在結構設計層面而工程編排層面已經標準化。理解了這一點你就可以把精力放到真正值得研究的地方。下一步建議從這個方向開始實踐先把模板在 CIFAR-10 上完整跑一遍重點觀察 ResNet 和 ViT 在相同數據下的行為差異然后用自己的數據集替換 CIFAR-10按實際分類數量修改build_model中的num_classes。如果數據量較大可以嘗試使用混合精度訓練如果模型較大可以考慮分布式訓練如果涉及目標檢測或語義分割則需要在本文模板基礎上增加對應任務頭。最后提醒一句模板是有適用邊界的。分類任務的訓練管線收斂了不代表所有視覺任務都適合這個套路。當你進入檢測、分割、多模態(tài)等方向時可以復用本文“注冊表管理模型、固定訓練循環(huán)、統一導出格式”的思維方式但代碼結構需要重新設計。把這份模板收藏起來作為你快速驗證模型和建立實驗基線的起點它會幫你省下大量重復造輪子的時間。 PyTorch統一視覺分類模板：從AlexNet到ViT訓練與ONNX部署
很多剛接觸視覺算法的同學都會遇到一個類似困境在文章或課程里把 AlexNet、VGG、ResNet、ViT 的結構圖看得明明白白真正打開編輯器準備訓練時卻發(fā)現自己只有“網絡結構代碼”沒有一套完整的數據加載、訓練循環(huán)、評估和導出腳本。于是從 A 倉庫復制數據增強從 B 倉庫找訓練參數再從 C 倉庫拼接模型調用最后訓練到一半才發(fā)現學習率策略和優(yōu)化器根本不匹配。這不是個人能力問題而是當前視覺入門教程普遍存在的空白結構講解很多工程閉環(huán)太少?？梢越o出一個明確判斷在圖像分類這個方向上從經典 CNN 到 Transformer訓練管線其實早已收斂。數據加載、損失函數、優(yōu)化器、評估邏輯、模型導出這些環(huán)節(jié)幾乎是固定的真正變化的只有模型自身以及少數幾處超參數。因此完全可以用一份統一代碼模板在 AlexNet、VGG、ResNet、ViT 四個主流架構之間橫向切換。這篇文章會圍繞一個可運行的分類項目展開用 PyTorch 實現一個統一模型構建層用同一套訓練腳本跑通四種架構再把訓練好的模型導出為 ONNX 做部署驗證。讀完你會擁有一份可以復用的視覺分類模板后續(xù)換成自己的數據和模型結構只需要改少量代碼。1. 這篇文章真正要解決的問題在進入代碼之前先把問題說清楚。很多人學習視覺模型時傾向于“每個模型單獨學一遍”但這樣會陷入一個誤區(qū)把大量的時間和注意力放在結構代碼的差異上而忽略了訓練流程的共性。AlexNet 和 ViT 雖然看起來差別巨大但它們都可以塞進一個標準的監(jiān)督分類流程里讀取圖像、前向傳播、計算交叉熵損失、反向傳播、更新參數、周期性評估準確率。這個模板要解決的三個核心問題是如何不寫四份幾乎一模一樣的訓練腳本只維護一份訓練代碼就完成多架構實驗。如何在四種架構之間做橫向對比而不是每次都在不同的代碼風格里來回切換。如何在訓練完成后把模型導出為通用格式ONNX讓部署代碼和訓練代碼解耦。對讀者來說這篇文章的技術收益是掌握一個“模型層可插拔、訓練層固定、導出層統一”的最小工程范式。如果你是在校學生可以用它快速完成模型對比實驗如果你是算法工程師可以用它作為新項目的主干腳本如果你剛入門深度學習可以通過這份模板搞清楚一條完整訓練鏈路到底由哪些環(huán)節(jié)組成。這里要先做一個提醒模板的價值在于“統一”但不等于“無腦”。四個架構在超參數敏感度上仍然有明顯差異這一點我會在后面的章節(jié)專門說明。代碼統一解決的是工程重復問題調參經驗解決的是模型差異問題兩者并不沖突。2. 四個架構并不需要四套訓練代碼先把四個架構放在一起對比這是理解整篇文章的基礎。模型提出年份核心思想典型結構更適用的場景AlexNet2012深度卷積 ReLU Dropout5 層卷積 3 層全連接學習基線、新手入門VGG2014小卷積核反復堆疊多個卷積塊 全連接結構直觀、遷移學習基線ResNet2015殘差連接緩解網絡退化基本塊 / 瓶頸塊通用視覺任務首選ViT2020圖像切塊 自注意力Patch Embedding Transformer Encoder大數據 預訓練微調AlexNet 的歷史意義在于它證明了“深度”的價值更深的網絡、非線性的 ReLU 激活、數據增強和 Dropout 共同把 ImageNet 分類準確率提升到了一個此前難以想象的高度。雖然今天直接在生產環(huán)境中使用 AlexNet 的團隊已經不多但它仍然是理解卷積神經網絡基本組成的最佳起點。VGG 的核心貢獻是簡化了卷積核的設計全部使用 3x3 小卷積核通過堆疊卷積塊來增加深度。它把“如何設計卷積層”變成了“如何堆疊卷基層”這個思路對后來很多網絡都有影響。缺點也明顯參數多、計算量大訓練和推理成本偏高。ResNet 解決了深度網絡難以訓練的問題。網絡太深時準確率反而會下降這叫“退化問題”。ResNet 通過殘差連接讓每一層學習的是“殘差”——也就是輸入和輸出之間的差值。這樣梯度可以繞過大部分卷積層直接向后傳播網絡可以做到很深而不容易訓練失敗。在實際工程中ResNet 系列往往是默認的視覺基線模型。ViT 則完全改變思路把圖像切成固定大小的 patch映射成 token然后交給 Transformer 的自注意力機制去建模全局關系。ViT 在超大數據集上訓練后再微調效果可以超過 CNN但在數據量不足時它不如 CNN 穩(wěn)。理解 ViT 的核心不需要寫完整代碼因為主流框架已經提供官方實現你需要做的是把它接入統一訓練流程。從工程視角看這四類架構的差異集中在“特征提取器”和“分類頭結構”上而訓練上下文完全一致。真正應該掌握的判斷是學習重點要分成兩層——結構層看差異工程層看共性。3. 環(huá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 預訓練模型在 torchvision 中從 0.13 開始提供。onnx 和 onnxruntime用于導出和部署驗證。安裝命令pip install torch torchvision onnx onnxruntime tqdm如果你需要 GPU 加速請根據自己機器的 CUDA 版本選擇對應的 PyTorch 安裝源這一點以 PyTorch 官網給出的安裝命令為準。CPU 環(huán)境也能運行本文全部代碼只是訓練速度會明顯變慢建議先用小 batch size 做驗證。版本問題是最容易踩坑的環(huán)節(jié)。torchvision 在較新版本中把pretrainedTrue改成了weights...的形式本文代碼為了清晰會直接使用新版寫法。如果你使用的是舊版本需要把代碼中的weightsIMAGENET1K_V1換成pretrainedTrue。這個細節(jié)會在常見問題章節(jié)繼續(xù)展開。4. 數據準備先跑通 CIFAR-10本文示例使用 CIFAR-10 數據集。它是視覺分類最常用的教學數據集50000 張訓練圖片10000 張測試圖片共 10 個類別。數據體積小、下載快特別適合用來驗證模板是否跑通。需要注意一個兩難問題CIFAR-10 原始圖片只有 32x32而 AlexNet、VGG、ResNet、ViT 的標準輸入通常是 224x224。直接把小圖 resize 到大尺寸會損失細節(jié)也會增加顯存占用但如果完全不 resize某些模型結構可能無法正常工作尤其是依賴固定 patch 數量的 ViT。模板的處理方式很直接統一 resize 到 224x224。這樣四個模型接受完全相同的輸入代碼路徑是統一的符合本文“模板化”的目標。需要理解的是從零訓練時把小圖放大到 224 并不理想所以當你把模板用到真實數據集時應該根據數據原始尺寸調整Resize的數值。下面是數據加載部分的代碼。這里只寫一份訓練集和驗證集都使用同一套邏輯區(qū)別在于是否做隨機數據增強。# 文件路徑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, ])訓練階段加入隨機水平翻轉和顏色抖動是為了提高模型泛化能力。驗證階段不做隨機增強只做 resize 和歸一化保證評估結果的穩(wěn)定性。在train.py中加載數據時使用torchvision.datasets.CIFAR10并設置downloadTrue。第一次運行會自動下載數據之后會直接讀取本地緩存。數據規(guī)模較小所以 DataLoader 的num_workers設為 2 即可過大會在 Windows 上引發(fā)額外問題。5. 統一模型構建層用注冊表管理四種架構很多人會問為什么不用簡單的 if-else 來切換模型if-else 當然能跑但隨著模型種類增多代碼會越來越混亂。更好的做法是用一個小的注冊表模式每個模型由一個 builder 函數負責構建函數名和模型字符串綁定。以后新增模型只需要加一個函數不需要改調用方。下面這段代碼是模板的關鍵部分。它封裝了四個模型的構建邏輯并且支持從 ImageNet 預訓練權重初始化。# 文件路徑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這段代碼里最重要的一層含義是模型構建是“可插拔”的。調用方永遠只依賴build_model(name, num_classes, pretrained)至于 name 對應的是卷積網絡還是 Transformer調用方不需要關心。這樣訓練腳本可以徹底固定下來不會因為換了模型而改動主流程。如果你使用的是較老版本的 torchvisionweights參數不存在。更保守的做法是把weights改成pretrained但那樣會失去新版本 API 的優(yōu)勢。這里建議優(yōu)先升級 torchvision而不是為了兼容舊版本把代碼寫得更復雜。6. 訓練循環(huán)一份腳本跑通四個模型訓練主腳本是整個模板的核心。它做的事情是所有分類任務共用的加載數據、構建模型、定義損失函數和優(yōu)化器、迭代若干輪、每輪驗證并保存最佳權重。下面給出完整訓練代碼。這是一個最小可運行版本沒有引入分布式訓練和復雜的日志框架方便理解主干邏輯。# 文件路徑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訓練完成最佳驗證準確率{best_acc * 100:.2f}%) print(f權重已保存至{os.path.join(args.save_dir, f{args.model}_best.pt)}) if __name__ __main__: main()這段代碼的邏輯并不復雜但有三個地方需要重點看優(yōu)化器支持 SGD 和 AdamW。ResNet、VGG 使用帶動量的 SGD 通常更穩(wěn)ViT 更推薦 AdamW配合較小學習率。學習率調度使用余弦退火。它能在訓練后期逐步降低學習率提升收斂穩(wěn)定性。每輪結束后立即在驗證集上評估并保存驗證準確率最高的權重。這樣訓練中斷時不至于丟失最佳結果。如果你想繼續(xù)跑第二個模型不需要改動任何代碼只需換--model參數。訓練腳本本身不感知模型內部結構差異這正是模板化的價值所在。7. 運行與驗證從命令行到日志判斷在開始訓練之前先用最小的配置驗證模板本身是否工作正常。這里建議先把 batch size 調小比如 16確定數據加載和模型 forward 沒有問題再跑正式訓練。四個模型分別通過以下命令啟動# AlexNet從零訓練 python train.py --model alexnet --epochs 30 --batch-size 64 --lr 1e-2 # VGG13從零訓練 python train.py --model vgg13 --epochs 30 --batch-size 64 --lr 1e-2 # ResNet18從零訓練 python train.py --model resnet18 --epochs 30 --batch-size 64 --lr 1e-2 # ViT使用 AdamW 并且加載 ImageNet 預訓練權重推薦 python train.py --model vit_b_16 --epochs 30 --batch-size 64 --lr 1e-4 --optimizer adamw --pretrained如果你機器顯存有限可以繼續(xù)降低 batch size。Batch size 不影響代碼結構只影響訓練速度和穩(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如何判斷訓練是否健康先看 train loss 是否持續(xù)下降。如果 train loss 下降但 val acc 不漲說明過擬合如果 train loss 和 val acc 都不動大概率是學習率設置有問題或者數據加載異常。不同模型在這個模板上的表現會遵循一些規(guī)律。ResNet18 從零訓練到 30 輪驗證準確率通常能接近 80% 甚至更高比較穩(wěn)定。AlexNet 參數多但結構相對淺在 CIFAR-10 這種小圖上提升較慢。VGG13 訓練速度慢顯存占用高但準確率通常不錯。ViT 如果不做預訓練在 CIFAR-10 上很容易過擬合這也是為什么它在小數據集上口碑不如 CNN 的直接原因。穿針引線地說模板讓你快速得到結果但理解不同模型的行為差異仍然要靠實驗對比和調參經驗。8. 導出 ONNX訓練與部署的解耦模型訓練完成后下一步通常是部署。很多人直接把 PyTorch 模型放到服務里這樣雖然能用但推理代碼和訓練代碼耦合很重。更通用的做法是導出為 ONNX再用 ONNX Runtime 這種專門的推理引擎加載。ONNX 的好處是與框架無關后續(xù)在 CPU、GPU 甚至邊緣設備上都能找到對應的運行環(huán)境。下面這段代碼把訓練好的 checkpoint 導出為 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 導出完成{args.output}) if __name__ __main__: main()導出命令python export_onnx.py --model resnet18 --weight checkpoints/resnet18_best.pt --output resnet18.onnx導出 ONNX 時最關鍵的是dynamic_axes。如果不設置ONNX 模型輸入 batch 會被固定為 1部署時只能一次推理一張圖。設置dynamic_axes后batch 維度可以是動態(tài)的更加靈活。導出完成后用 ONNX Runtime 做一次推理驗證確認導出沒有損壞模型結構。# 文件路徑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(預測類別索引:, logits.argmax(axis1))這一步的輸出應該是一個長度為 10 的向量因為 CIFAR-10 有 10 個類別。如果你導出的模型分類頭是 1000 類而不是 10 類說明導出腳本構建模型時沒有替換分類頭或者加載的權重和模型結構不匹配。這里也能看出模板設計的另一個收益訓練權重文件里保存的是“模型參數”而不是完整模型對象導出腳本只依賴build_model這個名字來重建結構。所以模型定義層必須保持穩(wěn)定這也是前面使用注冊表模式的原因。9. 常見問題與排查思路在實際運行中讀者最常遇到的并不是模型理論問題而是環(huán)境、版本、顯存這類工程問題。下面整理了一份排查表按優(yōu)先級排列。問題現象可能原因排查方式解決方案導入vit_b_16報 ImportErrortorchvision 版本過低沒有 ViT 實現執(zhí)行python -c import torchvision; print(torchvision.__version__)升級 torchvision 到 0.13 以上weights參數不被識別代碼使用了新版 API舊版只支持pretrainedTrue查看安裝的 PyTorch 版本對應文檔改為pretrainedTrue或升級版本顯存溢出OOMbatch size 過大或 224x224 輸入在 ViT 上顯存消耗高觀察日志中哪個階段報錯降低 batch size 重試--batch-size 16或更小并減少num_workerstrain loss 不下降學習率過大或過小也可能是數據沒有正確歸一化打日志檢查 loss 數量級嘗試 lr 1e-3從零訓練用 SGD lr1e-2預訓練微調用 lr1e-3ViT 在 CIFAR-10 上效果差小數據量 無預訓練Transformer 容易過擬合觀察 val acc 是否遠低于 train acc使用--pretrained或換用 ResNet 作為基線ONNX 導出后 batch 固定為 1dynamic_axes沒有正確設置用 Netron 打開 ONNX 文件查看輸入 shape在torch.onnx.export中設置dynamic_axes加載權重時 shape 不匹配checkpoint 的模型和當前--model不一致對比訓練命令和導出命令中的模型名確保訓練和導出使用同一個模型名稱驗證準確率高于訓練準確率Dropout 和 BatchNorm 在 train/eval 模式行為不同檢查model.train()和model.eval()是否正確切換在推理和驗證前必須執(zhí)行model.eval()這張表里的多數問題都無法靠“重裝環(huán)境”解決而是源于對版本差異和模型行為的不熟悉。建議你在開始訓練前先確認三件事torchvision 版本足夠新數據集能夠正常下載CIFAR-10 的歸一化參數已經配置正確。10. 工程實踐建議與超參選擇經驗模板能讓你快速把所有模型跑起來但在真實項目中使用時還需要遵守一些工程經驗否則很容易被表面上的“跑通”欺騙。第一先小后大。任何新的數據集接入模板都不要直接全量訓練。先用每個類別 100 到 200 張圖的子集跑 1 到 2 個 epoch確認數據加載正確、loss 能下降、模型可以收斂到隨機水平以上。工程里大量 bug 都是在全量訓練跑了幾小時后才暴露的小規(guī)模驗證成本低得多。第二學習率策略要區(qū)分模型。ResNet 和 VGG 使用 SGD momentum 是經過大量實踐驗證的穩(wěn)妥選擇學習率從 1e-2 開始調。ViT 更推薦 AdamW配合 warmup 和較小的學習率比如 1e-4 到 5e-4。如果套用同一個學習率到所有模型不是模板的錯而是沒有理解模型差異的體現。第三數據處理比模型結構更影響最終結果。CIFAR-10 因為圖小簡單增強就夠用但真實數據往往需要做 resize、隨機裁剪、旋轉、亮度擾動等多種組合。當模板遷移到新數據時首先要調整的是get_transform而不是急著換更復雜的模型。第四checkpoint 里不僅保存模型權重還可以保存優(yōu)化器狀態(tài)、當前 epoch、學習率等信息。本文示例為了簡潔只保存了模型權重和驗證準確率工程中建議至少加上優(yōu)化器狀態(tài)便于中斷后從最近的狀態(tài)繼續(xù)訓練。第五從零訓練還是微調預訓練權重必須想清楚。大數據量任務可以從零訓練但視覺任務普遍推薦使用 ImageNet 預訓練權重做初始化尤其是 ViT。預訓練初始化相當于讓模型從“已經會看邊的表示”開始而不是從“純隨機”開始這在中小數據集上差別非常大。第六部署前要檢查eval()狀態(tài)。PyTorch 模型在導出 ONNX 前一般要求調用model.eval()因為 BatchNorm 和 Dropout 在訓練與推理時的行為不同。如果不做導出的模型行為可能不穩(wěn)定。11. 總結與后續(xù)學習方向這份模板的核心價值不在于代碼技巧而在于把視覺分類任務的工程鏈路固化下來數據、模型、訓練、評估、導出。用同一份代碼跑 AlexNet、VGG、ResNet、ViT本質上是承認一個事實——這些模型的差異在結構設計層面而工程編排層面已經標準化。理解了這一點你就可以把精力放到真正值得研究的地方。下一步建議從這個方向開始實踐先把模板在 CIFAR-10 上完整跑一遍重點觀察 ResNet 和 ViT 在相同數據下的行為差異然后用自己的數據集替換 CIFAR-10按實際分類數量修改build_model中的num_classes。如果數據量較大可以嘗試使用混合精度訓練如果模型較大可以考慮分布式訓練如果涉及目標檢測或語義分割則需要在本文模板基礎上增加對應任務頭。最后提醒一句模板是有適用邊界的。分類任務的訓練管線收斂了不代表所有視覺任務都適合這個套路。當你進入檢測、分割、多模態(tài)等方向時可以復用本文“注冊表管理模型、固定訓練循環(huán)、統一導出格式”的思維方式但代碼結構需要重新設計。把這份模板收藏起來作為你快速驗證模型和建立實驗基線的起點它會幫你省下大量重復造輪子的時間。