別詳解:PyTorch實(shí)現(xiàn)與網(wǎng)絡(luò)改造指南)
1. 從一次模型改造翻車說起前陣子幫一個(gè)做圖像分割的朋友改網(wǎng)絡(luò)結(jié)構(gòu)他拿著一個(gè)分類網(wǎng)絡(luò)想直接套到分割任務(wù)上結(jié)果訓(xùn)練時(shí)loss死活降不下去推理出來的mask全是糊的。我一看代碼問題出在他把最后的全連接層刪掉之后直接接了個(gè)上采樣中間的特征圖維度完全對(duì)不上。這個(gè)坑其實(shí)特別典型——很多人學(xué)深度學(xué)習(xí)時(shí)卷積層和全連接層都寫過但真到要替換、要改造的時(shí)候腦子里對(duì)這兩者的本質(zhì)區(qū)別是模糊的。全連接層和全卷積層名字聽著像親戚實(shí)際在計(jì)算方式、參數(shù)結(jié)構(gòu)、適用場(chǎng)景上差別很大。這篇文章就從PyTorch代碼出發(fā)把這兩者的區(qū)別徹底講清楚。不管你是剛學(xué)完CNN基礎(chǔ)想進(jìn)一步理解網(wǎng)絡(luò)設(shè)計(jì)還是正在做分類到分割、檢測(cè)的遷移改造這些內(nèi)容都能直接拿來參考。我會(huì)把每一層的張量形狀變化、參數(shù)量計(jì)算、代碼實(shí)現(xiàn)都拆開講保證你看完能自己動(dòng)手改網(wǎng)絡(luò)。核心關(guān)鍵詞就三個(gè)全連接層、全卷積層、PyTorch實(shí)現(xiàn)。這三個(gè)詞貫穿全文也是理解兩者區(qū)別的關(guān)鍵抓手。2. 全連接層與全卷積層的本質(zhì)差異2.1 從數(shù)學(xué)運(yùn)算看兩者的根本不同全連接層的核心運(yùn)算就是矩陣乘法。假設(shè)輸入是一個(gè)長度為N的向量輸出是長度為M的向量那么這一層做的就是 y Wx b其中W是M×N的權(quán)重矩陣。每個(gè)輸出神經(jīng)元都和所有輸入神經(jīng)元相連所以叫“全連接”。在PyTorch里對(duì)應(yīng)的是nn.Linear(in_features, out_features)。全卷積層本質(zhì)上還是卷積操作但它和普通卷積在使用方式上有區(qū)別。普通卷積層通常后面會(huì)接全連接層做分類而全卷積層的意思是整個(gè)網(wǎng)絡(luò)從頭到尾只有卷積、池化、上采樣這類操作沒有任何全連接層。它輸出的是一個(gè)空間維度的特征圖而不是一個(gè)固定長度的向量。用一個(gè)生活化的類比全連接層像是一個(gè)“投票系統(tǒng)”每個(gè)輸入都要參與每個(gè)輸出的決策權(quán)重矩陣就是投票規(guī)則全卷積層像是一個(gè)“掃描儀”用一個(gè)固定大小的窗口在輸入上滑動(dòng)每個(gè)位置都做同樣的計(jì)算輸出的是同樣空間結(jié)構(gòu)的掃描結(jié)果。這個(gè)差異帶來的直接后果是全連接層要求輸入尺寸固定因?yàn)闄?quán)重矩陣的維度是寫死的全卷積層對(duì)輸入尺寸沒有硬性要求因?yàn)榫矸e核的參數(shù)量和輸入空間尺寸無關(guān)只和通道數(shù)、核大小有關(guān)。2.2 參數(shù)量對(duì)比為什么全連接層容易爆參數(shù)拿一個(gè)具體例子算一下。假設(shè)經(jīng)過前面的卷積和池化特征圖尺寸是7×7×512現(xiàn)在要接一個(gè)輸出10類的分類頭。如果用全連接層nn.Linear(7*7*512, 10)參數(shù)量是 7×7×512×10 10 250,890。這還只是7×7的情況如果是14×14×512參數(shù)量直接變成 14×14×512×10 10 ≈ 100萬。這就是為什么VGG這類網(wǎng)絡(luò)大部分參數(shù)都集中在最后幾個(gè)全連接層。如果用全卷積的方式假設(shè)用nn.Conv2d(512, 10, kernel_size7)參數(shù)量是 512×10×7×7 10 250,890。誒這個(gè)例子里參數(shù)量一樣對(duì)因?yàn)?×7的卷積核作用在7×7的輸入上感受野覆蓋了整個(gè)空間計(jì)算上等價(jià)于全連接。但關(guān)鍵在于如果輸入變成14×14全卷積層的參數(shù)量還是250,890而全連接層就爆了。注意全卷積層不是“沒有參數(shù)”而是參數(shù)量和輸入空間尺寸解耦。這是它最核心的優(yōu)勢(shì)。2.3 空間信息的保留與破壞全連接層會(huì)把空間維度展平flatten7×7×512變成25088的一維向量。這個(gè)操作直接丟掉了“哪個(gè)特征在哪個(gè)位置”的信息。對(duì)于分類任務(wù)我們只關(guān)心“圖里有沒有貓”不關(guān)心貓?jiān)谧笊辖沁€是右下角所以丟掉空間信息無所謂。但對(duì)于分割、檢測(cè)、關(guān)鍵點(diǎn)定位這類任務(wù)空間位置就是答案本身。你不可能把每個(gè)像素的類別都展平成一個(gè)向量再分類那樣參數(shù)量天文數(shù)字不說位置信息也全亂了。全卷積層輸出的特征圖保持了H×W的空間結(jié)構(gòu)每個(gè)位置都有對(duì)應(yīng)的預(yù)測(cè)結(jié)果天然適合這類密集預(yù)測(cè)任務(wù)。3. PyTorch代碼逐行拆解3.1 全連接層的PyTorch實(shí)現(xiàn)與形狀追蹤先寫一個(gè)最簡的全連接層示例把每一步的張量形狀都打印出來。import torch import torch.nn as nn # 模擬一個(gè)batch的數(shù)據(jù)batch_size2, 特征維度8 x torch.randn(2, 8) print(輸入形狀:, x.shape) # torch.Size([2, 8]) # 全連接層8維輸入4維輸出 fc nn.Linear(8, 4) y fc(x) print(輸出形狀:, y.shape) # torch.Size([2, 4]) # 查看權(quán)重和偏置的形狀 print(權(quán)重形狀:, fc.weight.shape) # torch.Size([4, 8]) print(偏置形狀:, fc.bias.shape) # torch.Size([4])這里nn.Linear(8, 4)的權(quán)重矩陣是4×8計(jì)算時(shí)是 x W.T b。輸入是(2, 8)輸出是(2, 4)。batch維度不受影響只變換特征維度。再看一個(gè)卷積網(wǎng)絡(luò)里常見的全連接層接法# 模擬卷積特征圖輸出batch2, channels512, H7, W7 feature torch.randn(2, 512, 7, 7) # 先展平 flattened feature.view(feature.size(0), -1) print(展平后形狀:, flattened.shape) # torch.Size([2, 25088]) # 接全連接層 fc nn.Linear(512 * 7 * 7, 10) output fc(flattened) print(分類輸出形狀:, output.shape) # torch.Size([2, 10])關(guān)鍵點(diǎn)在于view(feature.size(0), -1)這一步它把每個(gè)樣本的512×7×725088個(gè)值拉成一維。如果輸入圖片尺寸變了比如變成224×224經(jīng)過同樣的卷積后特征圖可能是7×7但如果輸入是320×320特征圖可能變成10×10這時(shí)候nn.Linear(512*7*7, 10)就會(huì)報(bào)維度不匹配的錯(cuò)誤。3.2 全卷積層的PyTorch實(shí)現(xiàn)與形狀追蹤同樣的輸入用全卷積的方式處理# 同樣的特征圖batch2, channels512, H7, W7 feature torch.randn(2, 512, 7, 7) # 用1x1卷積做通道變換輸出10個(gè)通道 conv1x1 nn.Conv2d(512, 10, kernel_size1) output conv1x1(feature) print(1x1卷積輸出形狀:, output.shape) # torch.Size([2, 10, 7, 7]) # 用7x7卷積輸出空間尺寸變成1x1 conv7x7 nn.Conv2d(512, 10, kernel_size7) output2 conv7x7(feature) print(7x7卷積輸出形狀:, output2.shape) # torch.Size([2, 10, 1, 1])注意nn.Conv2d(512, 10, kernel_size7)的輸出是(2, 10, 1, 1)如果把它view(2, 10)就和全連接層的輸出完全一樣。事實(shí)上當(dāng)卷積核大小等于輸入空間尺寸時(shí)卷積操作在數(shù)學(xué)上等價(jià)于全連接。但如果輸入變成14×14feature_large torch.randn(2, 512, 14, 14) output3 conv7x7(feature_large) print(大輸入7x7卷積輸出:, output3.shape) # torch.Size([2, 10, 8, 8])看到了嗎同樣的卷積層輸入變大輸出也變大參數(shù)量不變。這就是全卷積層對(duì)輸入尺寸魯棒的原因。3.3 兩者互相轉(zhuǎn)換的代碼實(shí)操把全連接層轉(zhuǎn)成等效的卷積層是網(wǎng)絡(luò)改造中的常見操作。假設(shè)有一個(gè)訓(xùn)練好的分類網(wǎng)絡(luò)最后的全連接層是nn.Linear(512*7*7, 10)可以這樣轉(zhuǎn)換# 原始全連接層 fc nn.Linear(512 * 7 * 7, 10) # 等效的卷積層 conv_equivalent nn.Conv2d(512, 10, kernel_size7) # 把全連接層的權(quán)重復(fù)制到卷積層 # fc.weight形狀是(10, 512*7*7)需要reshape成(10, 512, 7, 7) conv_equivalent.weight.data fc.weight.data.view(10, 512, 7, 7) conv_equivalent.bias.data fc.bias.data # 驗(yàn)證等價(jià)性 feature torch.randn(1, 512, 7, 7) flattened feature.view(1, -1) out_fc fc(flattened) out_conv conv_equivalent(feature).view(1, -1) print(全連接輸出:, out_fc) print(卷積輸出:, out_conv) print(最大差異:, (out_fc - out_conv).abs().max().item())實(shí)測(cè)下來兩者輸出差異在浮點(diǎn)誤差范圍內(nèi)通常小于1e-5。這個(gè)轉(zhuǎn)換在把分類網(wǎng)絡(luò)改成全卷積網(wǎng)絡(luò)時(shí)非常有用可以直接復(fù)用預(yù)訓(xùn)練權(quán)重。反過來把全卷積層轉(zhuǎn)成全連接層也有場(chǎng)景比如你想把某個(gè)全卷積網(wǎng)絡(luò)的輸出接一個(gè)固定類別的分類頭。但這種情況比較少因?yàn)槿矸e本身就能做分類取全局平均池化后接全連接也行。4. 實(shí)際網(wǎng)絡(luò)改造中的選擇與踩坑4.1 分類網(wǎng)絡(luò)改分割網(wǎng)絡(luò)的全卷積化流程拿一個(gè)典型的分類網(wǎng)絡(luò)舉例假設(shè)它前面是若干卷積和池化最后是全局平均池化加全連接。要改成分割網(wǎng)絡(luò)步驟大致如下第一步去掉最后的全局平均池化和全連接層。第二步把剩下的卷積部分作為編碼器。第三步接一個(gè)解碼器用轉(zhuǎn)置卷積或插值上采樣逐步恢復(fù)空間分辨率。第四步最后用一個(gè)1×1卷積輸出類別數(shù)通道的特征圖。class SimpleSegNet(nn.Module): def __init__(self, num_classes): super().__init__() # 編碼器部分假設(shè)輸入3通道 self.encoder nn.Sequential( nn.Conv2d(3, 64, 3, padding1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(64, 128, 3, padding1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(128, 256, 3, padding1), nn.ReLU(), ) # 解碼器部分 self.decoder nn.Sequential( nn.ConvTranspose2d(256, 128, 2, stride2), nn.ReLU(), nn.ConvTranspose2d(128, 64, 2, stride2), nn.ReLU(), nn.Conv2d(64, num_classes, 1) # 1x1卷積輸出類別 ) def forward(self, x): feat self.encoder(x) out self.decoder(feat) return out model SimpleSegNet(num_classes5) test_input torch.randn(2, 3, 128, 128) test_output model(test_input) print(分割輸出形狀:, test_output.shape) # torch.Size([2, 5, 128, 128])注意最后的輸出是(2, 5, 128, 128)每個(gè)像素位置都有5個(gè)類別的預(yù)測(cè)值。如果要做像素級(jí)分類直接對(duì)這個(gè)輸出做softmax然后算交叉熵就行。4.2 輸入尺寸不固定時(shí)的處理策略全卷積網(wǎng)絡(luò)雖然對(duì)輸入尺寸魯棒但實(shí)際用的時(shí)候還是要注意幾個(gè)點(diǎn)。第一如果網(wǎng)絡(luò)里有下采樣操作比如stride2的卷積或池化輸入尺寸最好是2的冪次或者能被整除的數(shù)否則上采樣后尺寸可能對(duì)不上。第二如果用了padding要注意padding方式對(duì)邊界的影響。一個(gè)實(shí)用的技巧是在forward里動(dòng)態(tài)獲取輸入尺寸然后在上采樣時(shí)用nn.functional.interpolate指定目標(biāo)尺寸。import torch.nn.functional as F class FlexibleSegNet(nn.Module): def __init__(self, num_classes): super().__init__() self.conv1 nn.Conv2d(3, 64, 3, padding1) self.conv2 nn.Conv2d(64, 128, 3, padding1) self.conv3 nn.Conv2d(128, num_classes, 1) def forward(self, x): input_size x.shape[2:] # 記錄輸入的空間尺寸 x F.relu(self.conv1(x)) x F.max_pool2d(x, 2) x F.relu(self.conv2(x)) x F.max_pool2d(x, 2) x self.conv3(x) # 上采樣回輸入尺寸 x F.interpolate(x, sizeinput_size, modebilinear, align_cornersFalse) return x這樣不管輸入是128×128還是256×256輸出都和輸入空間尺寸一致。4.3 參數(shù)量與顯存占用的實(shí)測(cè)對(duì)比在同樣的特征圖尺寸下全連接層和全卷積層的參數(shù)量可能相同但顯存占用和計(jì)算方式有差異。全連接層需要先把特征圖展平這個(gè)操作本身不占額外顯存view是共享內(nèi)存的但全連接層的矩陣乘法在底層實(shí)現(xiàn)上可能不如卷積高效。實(shí)測(cè)一個(gè)具體場(chǎng)景輸入(8, 512, 7, 7)輸出10類。層類型配置參數(shù)量前向顯存(約)全連接Linear(51277, 10)250,890較大全卷積Conv2d(512, 10, 7)250,890較小全卷積Conv2d(512, 10, 1) 全局池化5,130最小1×1卷積加全局平均池化的方案參數(shù)量最少而且對(duì)輸入尺寸完全魯棒。這也是很多現(xiàn)代網(wǎng)絡(luò)如ResNet、DenseNet最后用全局平均池化代替全連接層的原因。實(shí)操心得如果你的任務(wù)不需要空間輸出只是做分類優(yōu)先考慮全局平均池化加1×1卷積或直接接線性層。如果要做密集預(yù)測(cè)全卷積是唯一選擇。5. 常見問題與排查技巧5.1 維度不匹配報(bào)錯(cuò)的排查思路最常見的報(bào)錯(cuò)就是RuntimeError: mat1 and mat2 shapes cannot be multiplied或者size mismatch。排查步驟第一在報(bào)錯(cuò)的那一行前面打印輸入張量的形狀。第二確認(rèn)全連接層的in_features是否等于輸入張量的最后一維。第三如果是卷積網(wǎng)絡(luò)接全連接確認(rèn)展平后的維度計(jì)算是否正確。一個(gè)容易犯的錯(cuò)誤是輸入圖片尺寸變了但全連接層的in_features還是按原來的尺寸寫的。比如訓(xùn)練時(shí)用224×224推理時(shí)用256×256卷積部分輸出的特征圖空間尺寸變了展平后維度就對(duì)不上。解決方法有兩種一是推理時(shí)也resize到224×224二是把全連接層換成全卷積層或全局平均池化。5.2 全卷積網(wǎng)絡(luò)輸出尺寸對(duì)不上的處理全卷積網(wǎng)絡(luò)里如果有下采樣和上采樣輸出尺寸可能和輸入差幾個(gè)像素。比如輸入129×129經(jīng)過兩次stride2的池化變成33×33向下取整再經(jīng)過兩次轉(zhuǎn)置卷積上采樣變成132×132和輸入129×129不一致。處理方式用F.interpolate強(qiáng)制對(duì)齊到輸入尺寸或者在網(wǎng)絡(luò)設(shè)計(jì)時(shí)確保輸入尺寸是2的冪次。如果做分割任務(wù)輸出和輸入尺寸不一致會(huì)導(dǎo)致loss計(jì)算報(bào)錯(cuò)所以這一步必須處理。5.3 預(yù)訓(xùn)練權(quán)重加載時(shí)的層名匹配問題把分類網(wǎng)絡(luò)的全連接層換成卷積層后加載預(yù)訓(xùn)練權(quán)重時(shí)卷積層的參數(shù)名和原來的全連接層不一樣。比如原來叫classifier.0.weight現(xiàn)在叫backbone.conv_out.weight。這時(shí)候需要手動(dòng)映射。# 假設(shè)pretrained_dict是原始模型的權(quán)重 pretrained_dict torch.load(model.pth) model_dict model.state_dict() # 手動(dòng)映射需要轉(zhuǎn)換的層 pretrained_dict[backbone.conv_out.weight] pretrained_dict.pop(classifier.0.weight).view(10, 512, 7, 7) pretrained_dict[backbone.conv_out.bias] pretrained_dict.pop(classifier.0.bias) # 過濾掉不需要的層 pretrained_dict {k: v for k, v in pretrained_dict.items() if k in model_dict} model_dict.update(pretrained_dict) model.load_state_dict(model_dict)這個(gè)操作在實(shí)際改造中很常見尤其是用ImageNet預(yù)訓(xùn)練模型做遷移學(xué)習(xí)時(shí)。5.4 常見問題速查表問題現(xiàn)象可能原因解決方法全連接層維度報(bào)錯(cuò)輸入尺寸變化導(dǎo)致展平維度改變改用全卷積或全局平均池化全卷積輸出尺寸不對(duì)下采樣和上采樣不對(duì)稱用interpolate對(duì)齊尺寸預(yù)訓(xùn)練權(quán)重加載失敗層名或形狀不匹配手動(dòng)映射權(quán)重并reshape顯存爆了全連接層參數(shù)量過大換全局平均池化1x1卷積分割邊緣效果差上采樣丟失細(xì)節(jié)加skip connection融合淺層特征6. 幾個(gè)容易忽略的細(xì)節(jié)全連接層和全卷積層的選擇本質(zhì)上是在“參數(shù)量效率”和“空間信息保留”之間做權(quán)衡。分類任務(wù)里空間信息是冗余的全連接層或全局池化都能用密集預(yù)測(cè)任務(wù)里空間信息就是標(biāo)簽必須用全卷積。還有一個(gè)細(xì)節(jié)全連接層的權(quán)重矩陣是稠密的每個(gè)輸入都參與每個(gè)輸出的計(jì)算。全卷積層的卷積核是稀疏連接的每個(gè)輸出只和局部輸入有關(guān)。這個(gè)差異導(dǎo)致全連接層更容易過擬合尤其是在參數(shù)量大的時(shí)候。全卷積層因?yàn)閰?shù)共享泛化能力通常更好。最后分享一個(gè)我在實(shí)際改造網(wǎng)絡(luò)時(shí)的習(xí)慣先把原網(wǎng)絡(luò)在驗(yàn)證集上跑一遍記錄輸出形狀和關(guān)鍵層的特征圖尺寸。改造完之后用同樣的輸入再跑一遍逐層對(duì)比形狀變化。這樣能快速定位是哪一層出了問題比盲目調(diào)試效率高很多。