pytorch中有为efficientnet专门写好的网络模型,写在efficientnet_pytorch模块中。
模块包含EfficientNet的op-for-op的pytorch实现,也实现了预训练模型和示例。
Install via pip:
pip install efficientnet_pytorch
Or install from source:
git clone https://github.com/lukemelas/EfficientNet-PyTorch
cd EfficientNet-Pytorch
pip install -e .
From PyPI:
pip install keras_efficientnets
From Master branch:
pip install git+https://github.com/titu1994/keras-efficientnets.git
OR
git clone https://github.com/titu1994/keras-efficientnets.git
cd keras-efficientnets
pip install .
from efficientnet_pytorch import EfficientNet
model = EfficientNet.from_name('efficientnet-b0')
from efficientnet_pytorch import EfficientNet
model = EfficientNet.from_pretrained('efficientnet-b0')
from efficientnet_pytorch import EfficientNet
model = EfficientNet.from_pretrained('efficientnet-b5')
print(model)
只修改网络的最后几层(原始层结构):
(_conv_head): Conv2dStaticSamePadding(
512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False
(static_padding): Identity()
)
(_bn1): BatchNorm2d(2048, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)
(_fc): Linear(in_features=2048, out_features=1000, bias=True)
修改_fc层(最后一层),将输出的分类数由1000改为45:
得到in_features
:
feature = model._fc.in_features
print(feature)
结果:
Loaded pretrained weights for efficientnet-b5
2048
修改最后一层:
from efficientnet_pytorch import EfficientNet
from torch import nn
model = EfficientNet.from_pretrained('efficientnet-b5')
feature = model._fc.in_features
model._fc = nn.Linear(in_features=feature,out_features=45,bias=True)
print(model)
结果:
(_conv_head): Conv2dStaticSamePadding(
512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False
(static_padding): Identity()
)
(_bn1): BatchNorm2d(2048, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)
(_fc): Linear(in_features=2048, out_features=45, bias=True)
或者:和上述方法一致
from efficientnet_pytorch import EfficientNet
model = EfficientNet.from_pretrained('efficientnet-b5')
model._fc.out_features = 45
print(model)
结果:
(_conv_head): Conv2dStaticSamePadding(
512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False
(static_padding): Identity()
)
(_bn1): BatchNorm2d(2048, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)
(_fc): Linear(in_features=2048, out_features=45, bias=True)