Pytorch | 從零構建EfficientNet對CIFAR10進行分類
- CIFAR10數據集
- EfficientNet
- 設計理念
- 網絡結構
- 性能特點
- 應用領域
- 發展和改進
- EfficientNet結構代碼詳解
- 結構代碼
- 代碼詳解
- MBConv 類
- 初始化方法
- 前向傳播 forward 方法
- EfficientNet 類
- 初始化方法
- 前向傳播 forward 方法
- 訓練過程和測試結果
- 代碼匯總
- efficientnet.py
- train.py
- test.py
前面文章我們構建了AlexNet、Vgg、GoogleNet、ResNet、MobileNet對CIFAR10進行分類:
Pytorch | 從零構建AlexNet對CIFAR10進行分類
Pytorch | 從零構建Vgg對CIFAR10進行分類
Pytorch | 從零構建GoogleNet對CIFAR10進行分類
Pytorch | 從零構建ResNet對CIFAR10進行分類
Pytorch | 從零構建MobileNet對CIFAR10進行分類
這篇文章我們來構建EfficientNet.
CIFAR10數據集
CIFAR-10數據集是由加拿大高級研究所(CIFAR)收集整理的用于圖像識別研究的常用數據集,基本信息如下:
- 數據規模:該數據集包含60,000張彩色圖像,分為10個不同的類別,每個類別有6,000張圖像。通常將其中50,000張作為訓練集,用于模型的訓練;10,000張作為測試集,用于評估模型的性能。
- 圖像尺寸:所有圖像的尺寸均為32×32像素,這相對較小的尺寸使得模型在處理該數據集時能夠相對快速地進行訓練和推理,但也增加了圖像分類的難度。
- 類別內容:涵蓋了飛機(plane)、汽車(car)、鳥(bird)、貓(cat)、鹿(deer)、狗(dog)、青蛙(frog)、馬(horse)、船(ship)、卡車(truck)這10個不同的類別,這些類別都是現實世界中常見的物體,具有一定的代表性。
下面是一些示例樣本:
EfficientNet
EfficientNet是由谷歌大腦團隊在2019年提出的一種高效的卷積神經網絡架構,在圖像分類等任務上展現出了卓越的性能和效率,以下是對它的詳細介紹:
設計理念
- 平衡模型的深度、寬度和分辨率:傳統的神經網絡在提升性能時,往往只是單純地增加網絡的深度、寬度或輸入圖像的分辨率,而EfficientNet則通過一種系統的方法,同時對這三個維度進行優化調整,以達到在計算資源有限的情況下,模型性能的最大化。
網絡結構
- MBConv模塊:EfficientNet的核心模塊是MBConv(Mobile Inverted Residual Bottleneck),它基于深度可分離卷積和倒置殘差結構。這種結構在減少計算量的同時,能夠有效提取圖像特征,提高模型的表示能力。
- Compound Scaling方法:使用該方法對網絡的深度、寬度和分辨率進行統一縮放。通過一個固定的縮放系數,同時調整這三個維度,使得模型在不同的計算資源限制下,都能保持較好的性能和效率平衡。
上圖中是EfficientNet-B0的結構.
性能特點
- 高效性:在相同的計算資源下,EfficientNet能夠取得比傳統網絡更好的性能。例如,與ResNet-50相比,EfficientNet-B0在ImageNet數據集上取得了相近的準確率,但參數量和計算量卻大大減少。
- 可擴展性:通過Compound Scaling方法,可以方便地調整模型的大小,以適應不同的應用場景和計算資源限制。從EfficientNet-B0到EfficientNet-B7,模型的復雜度逐漸增加,性能也相應提升,能夠滿足從移動端到服務器端的不同需求。
應用領域
- 圖像分類:在ImageNet等大規模圖像分類數據集上,EfficientNet取得了領先的性能,成為圖像分類任務的首選模型之一。
- 目標檢測:與Faster R-CNN等目標檢測框架結合,EfficientNet作為骨干網絡,能夠提高目標檢測的準確率和速度,在Pascal VOC、COCO等數據集上取得了不錯的效果。
- 語義分割:在語義分割任務中,EfficientNet也展現出了一定的優勢,通過與U-Net等分割網絡結合,能夠對圖像進行像素級的分類,在Cityscapes等數據集上有較好的表現。
發展和改進
- EfficientNet v2:在EfficientNet基礎上進行了進一步優化,主要改進包括改進了漸進式學習的方法,在訓練過程中逐漸增加圖像的分辨率,使得模型能夠更好地學習到不同尺度的特征,同時優化了網絡結構,提高了模型的訓練速度和性能。
- 其他改進:研究人員還在EfficientNet的基礎上,結合其他技術如注意力機制、知識蒸餾等,進一步提高模型的性能和泛化能力。
EfficientNet結構代碼詳解
結構代碼
import torch
import torch.nn as nn
import torch.nn.functional as Fclass MBConv(nn.Module):def __init__(self, in_channels, out_channels, expand_ratio, kernel_size, stride, padding, se_ratio=0.25):super(MBConv, self).__init__()self.stride = strideself.use_res_connect = (stride == 1 and in_channels == out_channels)# 擴展通道數(如果需要)expanded_channels = in_channels * expand_ratioself.expand_conv = nn.Conv2d(in_channels, expanded_channels, kernel_size=1, padding=0, bias=False)self.bn1 = nn.BatchNorm2d(expanded_channels)# 深度可分離卷積self.depthwise_conv = nn.Conv2d(expanded_channels, expanded_channels, kernel_size=kernel_size,stride=stride, padding=padding, groups=expanded_channels, bias=False)self.bn2 = nn.BatchNorm2d(expanded_channels)# 壓縮和激勵(SE)模塊(可選,根據se_ratio判斷是否添加)if se_ratio > 0:se_channels = int(in_channels * se_ratio)self.se = nn.Sequential(nn.AdaptiveAvgPool2d(1),nn.Conv2d(expanded_channels, se_channels, kernel_size=1),nn.ReLU(inplace=True),nn.Conv2d(se_channels, expanded_channels, kernel_size=1),nn.Sigmoid())else:self.se = None# 投影卷積,恢復到輸出通道數self.project_conv = nn.Conv2d(expanded_channels, out_channels, kernel_size=1, padding=0, bias=False)self.bn3 = nn.BatchNorm2d(out_channels)def forward(self, x):identity = x# 擴展通道數out = F.relu(self.bn1(self.expand_conv(x)))# 深度可分離卷積out = F.relu(self.bn2(self.depthwise_conv(out)))# SE模塊操作(如果存在)if self.se is not None:se_out = self.se(out)out = out * se_out# 投影卷積out = self.bn3(self.project_conv(out))# 殘差連接(如果滿足條件)if self.use_res_connect:out += identityreturn outclass EfficientNet(nn.Module):def __init__(self, num_classes, width_coefficient=1.0, depth_coefficient=1.0, dropout_rate=0.2):super(EfficientNet, self).__init__()self.stem_conv = nn.Conv2d(3, 32, kernel_size=3, stride=2, padding=1, bias=False)self.bn1 = nn.BatchNorm2d(32)mbconv_config = [# (in_channels, out_channels, expand_ratio, kernel_size, stride, padding)(32, 16, 1, 3, 1, 1),(16, 24, 6, 3, 2, 1),(24, 40, 6, 5, 2, 2),(40, 80, 6, 3, 2, 1),(80, 112, 6, 5, 1, 2),(112, 192, 6, 5, 2, 2),(192, 320, 6, 3, 1, 1)]# 根據深度系數調整每個MBConv模塊的重復次數,這里簡單地向下取整,你也可以根據實際情況采用更合理的方式repeat_counts = [max(1, int(depth_coefficient * 1)) for _ in mbconv_config]layers = []for i, config in enumerate(mbconv_config):in_channels, out_channels, expand_ratio, kernel_size, stride, padding = configfor _ in range(repeat_counts[i]):layers.append(MBConv(int(in_channels * width_coefficient),int(out_channels * width_coefficient),expand_ratio, kernel_size, stride, padding))self.mbconv_layers = nn.Sequential(*layers)self.last_conv = nn.Conv2d(int(320 * width_coefficient), 1280, kernel_size=1, bias=False)self.bn2 = nn.BatchNorm2d(1280)self.avgpool = nn.AdaptiveAvgPool2d(1)self.dropout = nn.Dropout(dropout_rate)self.fc = nn.Linear(1280, num_classes)def forward(self, x):out = F.relu(self.bn1(self.stem_conv(x)))out = self.mbconv_layers(out)out = F.relu(self.bn2(self.last_conv(out)))out = self.avgpool(out)out = out.view(out.size(0), -1)out = self.dropout(out)out = self.fc(out)return out
代碼詳解
以下是對上述EfficientNet
代碼的詳細解釋,代碼整體定義了EfficientNet
網絡結構,主要由MBConv
模塊堆疊以及一些常規的卷積、池化和全連接層構成,下面按照類和方法分別進行剖析:
MBConv 類
這是EfficientNet
中的核心模塊,實現了MBConv
(Mobile Inverted Residual Bottleneck Convolution)結構,其代碼如下:
class MBConv(nn.Module):def __init__(self, in_channels, out_channels, expand_ratio, kernel_size, stride, padding, se_ratio=0.25):super(MBConv, self).__init__()self.stride = strideself.use_res_connect = (stride == 1 and in_channels == out_channels)# 擴展通道數(如果需要)expanded_channels = in_channels * expand_ratioself.expand_conv = nn.Conv2d(in_channels, expanded_channels, kernel_size=1, padding=0, bias=False)self.bn1 = nn.BatchNorm2d(expanded_channels)# 深度可分離卷積self.depthwise_conv = nn.Conv2d(expanded_channels, expanded_channels, kernel_size=kernel_size,stride=stride, padding=padding, groups=expanded_channels, bias=False)self.bn2 = nn.BatchNorm2d(expanded_channels)# 壓縮和激勵(SE)模塊(可選,根據se_ratio判斷是否添加)if se_ratio > 0:se_channels = int(in_channels * se_ratio)self.se = nn.Sequential(nn.AdaptiveAvgPool2d(1),nn.Conv2d(expanded_channels, se_channels, kernel_size=1),nn.ReLU(inplace=True),nn.Conv2d(se_channels, expanded_channels, kernel_size=1),nn.Sigmoid())else:self.se = None# 投影卷積,恢復到輸出通道數self.project_conv = nn.Conv2d(expanded_channels, out_channels, kernel_size=1, padding=0, bias=False)self.bn3 = nn.BatchNorm2d(out_channels)def forward(self, x):identity = x# 擴展通道數out = F.relu(self.bn1(self.expand_conv(x)))# 深度可分離卷積out = F.relu(self.bn2(self.depthwise_conv(out)))# SE模塊操作(如果存在)if self.se is not None:se_out = self.se(out)out = out * se_out# 投影卷積out = self.bn3(self.project_conv(out))# 殘差連接(如果滿足條件)if self.use_res_connect:out += identityreturn out
初始化方法
- 參數說明:
in_channels
:輸入張量的通道數。out_channels
:輸出張量的通道數。expand_ratio
:用于確定擴展通道數時的比例系數,決定是否對輸入通道數進行擴展以及擴展的倍數。kernel_size
:卷積核的大小,用于深度可分離卷積等操作。stride
:卷積的步長,控制特征圖在卷積過程中的下采樣程度等。padding
:卷積操作時的填充大小,保證輸入輸出特征圖尺寸在特定要求下的一致性等。se_ratio
(可選,默認值為0.25):用于控制壓縮和激勵(SE)模塊中通道壓縮的比例,若為0則不添加SE模塊。
- 初始化操作:
- 首先保存傳入的
stride
參數,并根據stride
和輸入輸出通道數判斷是否使用殘差連接(use_res_connect
),只有當步長為1且輸入輸出通道數相等時才使用殘差連接,這符合殘差網絡的基本原理,有助于梯度傳播和特征融合。 - 根據
expand_ratio
計算擴展后的通道數expanded_channels
,并創建expand_conv
卷積層用于擴展通道數,同時搭配對應的bn1
批歸一化層,對擴展后的特征進行歸一化處理,有助于加速網絡收斂和提升模型穩定性。 - 定義
depthwise_conv
深度可分離卷積層,其分組數設置為expanded_channels
,意味著每個通道單獨進行卷積操作,這種方式可以在減少計算量的同時保持較好的特征提取能力,同時搭配bn2
批歸一化層。 - 根據
se_ratio
判斷是否添加壓縮和激勵(SE)模塊。如果se_ratio
大于0,則創建一個包含自適應平均池化、卷積、激活函數(ReLU)、卷積和Sigmoid激活的序列模塊se
,用于對特征進行通道維度上的重加權,增強模型對不同通道特征的關注度;若se_ratio
為0,則將se
設為None
。 - 最后創建
project_conv
投影卷積層用于將擴展和處理后的特征恢復到指定的輸出通道數,并搭配bn3
批歸一化層。
- 首先保存傳入的
前向傳播 forward 方法
- 首先將輸入張量
x
保存為identity
,用于后續可能的殘差連接。 - 通過
F.relu(self.bn1(self.expand_conv(x)))
對輸入進行通道擴展,并使用ReLU激活函數和批歸一化進行處理,得到擴展后的特征表示。 - 接著對擴展后的特征執行深度可分離卷積操作
F.relu(self.bn2(self.depthwise_conv(out)))
,同樣使用ReLU激活和批歸一化處理。 - 如果存在SE模塊(
self.se
不為None
),則將經過深度可分離卷積后的特征傳入SE模塊進行通道重加權,即se_out = self.se(out)
,然后將特征與重加權后的結果相乘out = out * se_out
。 - 通過
self.bn3(self.project_conv(out))
進行投影卷積操作,將特征恢復到輸出通道數,并進行批歸一化處理。 - 最后,如果滿足殘差連接條件(
self.use_res_connect
為True
),則將投影卷積后的特征與最初保存的輸入特征identity
相加,實現殘差連接,最終返回處理后的特征張量。
EfficientNet 類
這是整體的EfficientNet
網絡模型類,代碼如下:
class EfficientNet(nn.Module):def __init__(self, num_classes, width_coefficient=1.0, depth_coefficient=1.0, dropout_rate=0.2):super(EfficientNet, self).__init__()self.stem_conv = nn.Conv2d(3, 32, kernel_size=3, stride=2, padding=1, bias=False)self.bn1 = nn.BatchNorm2d(32)mbconv_config = [# (in_channels, out_channels, expand_ratio, kernel_size, stride, padding)(32, 16, 1, 3, 1, 1),(16, 24, 6, 3, 2, 1),(24, 40, 6, 5, 2, 2),(40, 80, 6, 3, 2, 1),(80, 112, 6, 5, 1, 2),(112, 192, 6, 5, 2, 2),(192, 320, 6, 3, 1, 1)]# 根據深度系數調整每個MBConv模塊的重復次數,這里簡單地向下取整,你也可以根據實際情況采用更合理的方式repeat_counts = [max(1, int(depth_coefficient * 1)) for _ in mbconv_config]layers = []for i, config in enumerate(mbconv_config):in_channels, out_channels, expand_ratio, kernel_size, stride, padding = configfor _ in range(repeat_counts[i]):layers.append(MBConv(int(in_channels * width_coefficient),int(out_channels * width_coefficient),expand_ratio, kernel_size, stride, padding))self.mbconv_layers = nn.Sequential(*layers)self.last_conv = nn.Conv2d(int(320 * width_coefficient), 1280, kernel_size=1, bias=False)self.bn2 = nn.BatchNorm2d(1280)self.avgpool = nn.AdaptiveAvgPool2d(1)self.dropout = nn.Dropout(dropout_rate)self.fc = nn.Linear(1280, num_classes)def forward(self, x):out = F.relu(self.bn1(self.stem_conv(x)))out = self.mbconv_layers(out)out = F.relu(self.bn2(self.last_conv(out)))out = self.avgpool(out)out = out.view(out.size(0), -1)out = self.dropout(out)out = self.fc(out)return out
初始化方法
- 參數說明:
num_classes
:最終分類任務的類別數量,用于確定全連接層的輸出維度。width_coefficient
(默認值為1.0):用于控制模型各層的通道數,實現對模型寬度的縮放調整。depth_coefficient
(默認值為1.0):用于控制模型中MBConv
模塊的重復次數,實現對模型深度的縮放調整。dropout_rate
(默認值為0.2):在全連接層之前使用的Dropout概率,用于防止過擬合。
- 初始化操作:
- 首先創建
stem_conv
卷積層,它將輸入的圖像數據(通常通道數為3,對應RGB圖像)進行初始的卷積操作,步長為2,起到下采樣的作用,同時不使用偏置(bias=False
),并搭配bn1
批歸一化層對卷積后的特征進行歸一化處理。 - 定義
mbconv_config
列表,其中每個元素是一個元組,包含了MBConv
模塊的各項配置參數,如輸入通道數、輸出通道數、擴展比例、卷積核大小、步長和填充等,這是構建MBConv
模塊的基礎配置信息。 - 根據
depth_coefficient
計算每個MBConv
模塊的重復次數,通過列表推導式repeat_counts = [max(1, int(depth_coefficient * 1)) for _ in mbconv_config]
實現,這里簡單地將每個配置對應的重復次數設置為與depth_coefficient
成比例(向下取整且保證至少重復1次),你可以根據更精細的設計規則來調整這個計算方式。 - 構建
self.mbconv_layers
,通過兩層嵌套循環實現。外層循環遍歷mbconv_config
配置列表,內層循環根據對應的重復次數來多次添加同一個MBConv
模塊實例到layers
列表中,最后將layers
列表轉換為nn.Sequential
類型的模塊,這樣就實現了根據depth_coefficient
對模型深度進行調整以及MBConv
模塊的堆疊搭建。 - 創建
last_conv
卷積層,它將經過MBConv
模塊處理后的特征進行進一步的卷積操作,將通道數轉換為1280,同樣不使用偏置,搭配bn2
批歸一化層。 - 定義
avgpool
自適應平均池化層,將特征圖轉換為固定大小(這里為1x1),方便后續全連接層處理。 - 創建
dropout
Dropout層,按照指定的dropout_rate
在全連接層之前進行隨機失活操作,防止過擬合。 - 最后定義
fc
全連接層,其輸入維度為1280(經過池化后的特征維度),輸出維度為num_classes
,用于最終的分類預測。
- 首先創建
前向傳播 forward 方法
- 首先將輸入
x
傳入stem_conv
卷積層進行初始卷積,然后通過F.relu(self.bn1(self.stem_conv(x)))
進行激活和批歸一化處理,得到初始的特征表示。 - 將初始特征傳入
self.mbconv_layers
,即經過一系列堆疊的MBConv
模塊進行特征提取和變換,充分挖掘圖像中的特征信息。 - 接著對經過
MBConv
模塊處理后的特征執行F.relu(self.bn2(self.last_conv(out)))
操作,進行最后的卷積以及激活、批歸一化處理。 - 使用
self.avgpool(out)
進行自適應平均池化,將特征圖尺寸變為1x1,實現特征的壓縮和固定維度表示。 - 通過
out = out.view(out.size(0), -1)
將池化后的特征張量展平為一維向量,方便全連接層處理,這里-1
表示自動根據張量元素總數和已知的批量大小維度(out.size(0)
)來推斷展平后的維度大小。 - 然后將展平后的特征傳入
self.dropout(out)
進行Dropout操作,隨機丟棄一部分神經元,防止過擬合。 - 最后將特征傳入
self.fc(out)
全連接層進行分類預測,得到最終的輸出結果,輸出的維度與設定的num_classes
一致,表示每個樣本屬于各個類別的預測概率(或得分等,具體取決于任務和后續處理),并返回該輸出結果。
訓練過程和測試結果
訓練過程損失函數變化曲線:
訓練過程準確率變化曲線:
測試結果:
代碼匯總
項目github地址
項目結構:
|--data
|--models|--__init__.py|-efficientnet.py|--...
|--results
|--weights
|--train.py
|--test.py
efficientnet.py
import torch
import torch.nn as nn
import torch.nn.functional as Fclass MBConv(nn.Module):def __init__(self, in_channels, out_channels, expand_ratio, kernel_size, stride, padding, se_ratio=0.25):super(MBConv, self).__init__()self.stride = strideself.use_res_connect = (stride == 1 and in_channels == out_channels)# 擴展通道數(如果需要)expanded_channels = in_channels * expand_ratioself.expand_conv = nn.Conv2d(in_channels, expanded_channels, kernel_size=1, padding=0, bias=False)self.bn1 = nn.BatchNorm2d(expanded_channels)# 深度可分離卷積self.depthwise_conv = nn.Conv2d(expanded_channels, expanded_channels, kernel_size=kernel_size,stride=stride, padding=padding, groups=expanded_channels, bias=False)self.bn2 = nn.BatchNorm2d(expanded_channels)# 壓縮和激勵(SE)模塊(可選,根據se_ratio判斷是否添加)if se_ratio > 0:se_channels = int(in_channels * se_ratio)self.se = nn.Sequential(nn.AdaptiveAvgPool2d(1),nn.Conv2d(expanded_channels, se_channels, kernel_size=1),nn.ReLU(inplace=True),nn.Conv2d(se_channels, expanded_channels, kernel_size=1),nn.Sigmoid())else:self.se = None# 投影卷積,恢復到輸出通道數self.project_conv = nn.Conv2d(expanded_channels, out_channels, kernel_size=1, padding=0, bias=False)self.bn3 = nn.BatchNorm2d(out_channels)def forward(self, x):identity = x# 擴展通道數out = F.relu(self.bn1(self.expand_conv(x)))# 深度可分離卷積out = F.relu(self.bn2(self.depthwise_conv(out)))# SE模塊操作(如果存在)if self.se is not None:se_out = self.se(out)out = out * se_out# 投影卷積out = self.bn3(self.project_conv(out))# 殘差連接(如果滿足條件)if self.use_res_connect:out += identityreturn outclass EfficientNet(nn.Module):def __init__(self, num_classes, width_coefficient=1.0, depth_coefficient=1.0, dropout_rate=0.2):super(EfficientNet, self).__init__()self.stem_conv = nn.Conv2d(3, 32, kernel_size=3, stride=2, padding=1, bias=False)self.bn1 = nn.BatchNorm2d(32)mbconv_config = [# (in_channels, out_channels, expand_ratio, kernel_size, stride, padding)(32, 16, 1, 3, 1, 1),(16, 24, 6, 3, 2, 1),(24, 40, 6, 5, 2, 2),(40, 80, 6, 3, 2, 1),(80, 112, 6, 5, 1, 2),(112, 192, 6, 5, 2, 2),(192, 320, 6, 3, 1, 1)]# 根據深度系數調整每個MBConv模塊的重復次數,這里簡單地向下取整,你也可以根據實際情況采用更合理的方式repeat_counts = [max(1, int(depth_coefficient * 1)) for _ in mbconv_config]layers = []for i, config in enumerate(mbconv_config):in_channels, out_channels, expand_ratio, kernel_size, stride, padding = configfor _ in range(repeat_counts[i]):layers.append(MBConv(int(in_channels * width_coefficient),int(out_channels * width_coefficient),expand_ratio, kernel_size, stride, padding))self.mbconv_layers = nn.Sequential(*layers)self.last_conv = nn.Conv2d(int(320 * width_coefficient), 1280, kernel_size=1, bias=False)self.bn2 = nn.BatchNorm2d(1280)self.avgpool = nn.AdaptiveAvgPool2d(1)self.dropout = nn.Dropout(dropout_rate)self.fc = nn.Linear(1280, num_classes)def forward(self, x):out = F.relu(self.bn1(self.stem_conv(x)))out = self.mbconv_layers(out)out = F.relu(self.bn2(self.last_conv(out)))out = self.avgpool(out)out = out.view(out.size(0), -1)out = self.dropout(out)out = self.fc(out)return out
train.py
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from models import *
import matplotlib.pyplot as pltimport ssl
ssl._create_default_https_context = ssl._create_unverified_context# 定義數據預處理操作
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.491, 0.482, 0.446), (0.247, 0.243, 0.261))])# 加載CIFAR10訓練集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,download=False, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128,shuffle=True, num_workers=2)# 定義設備(GPU優先,若可用)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")# 實例化模型
model_name = 'EfficientNet'
if model_name == 'AlexNet':model = AlexNet(num_classes=10).to(device)
elif model_name == 'Vgg_A':model = Vgg(cfg_vgg='A', num_classes=10).to(device)
elif model_name == 'Vgg_A-LRN':model = Vgg(cfg_vgg='A-LRN', num_classes=10).to(device)
elif model_name == 'Vgg_B':model = Vgg(cfg_vgg='B', num_classes=10).to(device)
elif model_name == 'Vgg_C':model = Vgg(cfg_vgg='C', num_classes=10).to(device)
elif model_name == 'Vgg_D':model = Vgg(cfg_vgg='D', num_classes=10).to(device)
elif model_name == 'Vgg_E':model = Vgg(cfg_vgg='E', num_classes=10).to(device)
elif model_name == 'GoogleNet':model = GoogleNet(num_classes=10).to(device)
elif model_name == 'ResNet18':model = ResNet18(num_classes=10).to(device)
elif model_name == 'ResNet34':model = ResNet34(num_classes=10).to(device)
elif model_name == 'ResNet50':model = ResNet50(num_classes=10).to(device)
elif model_name == 'ResNet101':model = ResNet101(num_classes=10).to(device)
elif model_name == 'ResNet152':model = ResNet152(num_classes=10).to(device)
elif model_name == 'MobileNet':model = MobileNet(num_classes=10).to(device)
elif model_name == 'EfficientNet':model = EfficientNet(num_classes=10).to(device)criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)# 訓練輪次
epochs = 15def train(model, trainloader, criterion, optimizer, device):model.train()running_loss = 0.0correct = 0total = 0for i, data in enumerate(trainloader, 0):inputs, labels = data[0].to(device), data[1].to(device)optimizer.zero_grad()outputs = model(inputs)loss = criterion(outputs, labels)loss.backward()optimizer.step()running_loss += loss.item()_, predicted = outputs.max(1)total += labels.size(0)correct += predicted.eq(labels).sum().item()epoch_loss = running_loss / len(trainloader)epoch_acc = 100. * correct / totalreturn epoch_loss, epoch_accif __name__ == "__main__":loss_history, acc_history = [], []for epoch in range(epochs):train_loss, train_acc = train(model, trainloader, criterion, optimizer, device)print(f'Epoch {epoch + 1}: Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}%')loss_history.append(train_loss)acc_history.append(train_acc)# 保存模型權重,每5輪次保存到weights文件夾下if (epoch + 1) % 5 == 0:torch.save(model.state_dict(), f'weights/{model_name}_epoch_{epoch + 1}.pth')# 繪制損失曲線plt.plot(range(1, epochs+1), loss_history, label='Loss', marker='o')plt.xlabel('Epoch')plt.ylabel('Loss')plt.title('Training Loss Curve')plt.legend()plt.savefig(f'results\\{model_name}_train_loss_curve.png')plt.close()# 繪制準確率曲線plt.plot(range(1, epochs+1), acc_history, label='Accuracy', marker='o')plt.xlabel('Epoch')plt.ylabel('Accuracy (%)')plt.title('Training Accuracy Curve')plt.legend()plt.savefig(f'results\\{model_name}_train_acc_curve.png')plt.close()
test.py
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from models import *import ssl
ssl._create_default_https_context = ssl._create_unverified_context
# 定義數據預處理操作
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.491, 0.482, 0.446), (0.247, 0.243, 0.261))])# 加載CIFAR10測試集
testset = torchvision.datasets.CIFAR10(root='./data', train=False,download=False, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=128,shuffle=False, num_workers=2)# 定義設備(GPU優先,若可用)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")# 實例化模型
model_name = 'EfficientNet'
if model_name == 'AlexNet':model = AlexNet(num_classes=10).to(device)
elif model_name == 'Vgg_A':model = Vgg(cfg_vgg='A', num_classes=10).to(device)
elif model_name == 'Vgg_A-LRN':model = Vgg(cfg_vgg='A-LRN', num_classes=10).to(device)
elif model_name == 'Vgg_B':model = Vgg(cfg_vgg='B', num_classes=10).to(device)
elif model_name == 'Vgg_C':model = Vgg(cfg_vgg='C', num_classes=10).to(device)
elif model_name == 'Vgg_D':model = Vgg(cfg_vgg='D', num_classes=10).to(device)
elif model_name == 'Vgg_E':model = Vgg(cfg_vgg='E', num_classes=10).to(device)
elif model_name == 'GoogleNet':model = GoogleNet(num_classes=10).to(device)
elif model_name == 'ResNet18':model = ResNet18(num_classes=10).to(device)
elif model_name == 'ResNet34':model = ResNet34(num_classes=10).to(device)
elif model_name == 'ResNet50':model = ResNet50(num_classes=10).to(device)
elif model_name == 'ResNet101':model = ResNet101(num_classes=10).to(device)
elif model_name == 'ResNet152':model = ResNet152(num_classes=10).to(device)
elif model_name == 'MobileNet':model = MobileNet(num_classes=10).to(device)
elif model_name == 'EfficientNet':model = EfficientNet(num_classes=10).to(device)criterion = nn.CrossEntropyLoss()# 加載模型權重
weights_path = f"weights/{model_name}_epoch_15.pth"
model.load_state_dict(torch.load(weights_path, map_location=device))def test(model, testloader, criterion, device):model.eval()running_loss = 0.0correct = 0total = 0with torch.no_grad():for data in testloader:inputs, labels = data[0].to(device), data[1].to(device)outputs = model(inputs)loss = criterion(outputs, labels)running_loss += loss.item()_, predicted = outputs.max(1)total += labels.size(0)correct += predicted.eq(labels).sum().item()epoch_loss = running_loss / len(testloader)epoch_acc = 100. * correct / totalreturn epoch_loss, epoch_accif __name__ == "__main__":test_loss, test_acc = test(model, testloader, criterion, device)print(f"================{model_name} Test================")print(f"Load Model Weights From: {weights_path}")print(f'Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.2f}%')