EfficientNet是目前图像分类中最好的网络之一了,参数数量小(较VGG和Inception都小好几倍),运行速度快,github上的地址为:https://github.com/lukemelas/EfficientNet-PyTorch。对论文的解读可以百度到很多,学习下理论可以,但是对模型的复现和ImageNet的训练,对于我只有2G显存显卡的笔记本电脑来说,还是别做这个梦了。
git下来源码,自己下载了efficientnet-b3模型参数文件(github上速度超慢,49M的文件花了我一个晚上才下完,满满都是泪。虽然码云上也有别人复制过来的仓库,但是全部都只复制了代码,没有作者预训练好的模型参数。),作者说可以直接pip install efficient-pytorch安装,但这样安装后发现有个问题:每次运行都要重新从github上下载模型文件,不能指定模型文件的路径加载。因此对源码做了一点修改,如下。
1、在源码 <EfficientNet-PyTorch-Root>/efficientnet-pytorch/目录下修改utils.py文件,添加自己的函数
def load_local_weights(model,filepath,load_fc=True):
""" Loads pretrained weights from local file ."""
if (not os.path.exists(filepath)):
raise ValueError('%s does not exist.' % filepath)
state_dict=torch.load(filepath)
if load_fc:
model.load_state_dict(state_dict)
else:
state_dict.pop('_fc.weight')
state_dict.pop('_fc.bias')
res = model.load_state_dict(state_dict, strict=False)
assert set(res.missing_keys) == set(['_fc.weight', '_fc.bias']), 'issue loading pretrained weights'
print('Loaded pretrained weights for {}'.format(filepath))
2、在源码 <EfficientNet-PyTorch-Root>/efficientnet-pytorch/目录下修改model.py文件,在EfficientNet类中添加如下函数、
@classmethod
def from_local_file(cls,model_name,model_file,num_classes=1000):
"""Loads weights parameters from local file"""
model=cls.from_name(model_name,override_params={'num_classes':num_classes})
load_local_weights(model,model_file, load_fc=(num_classes==1000))
return model
3、运行<EfficientNet-PyTorch-Root>/setup.py,重新构建包:
>>>cd <EfficientNet_Root>
>>>python setup.py build
>>>python setup.py install
就将包安装到我的虚拟环境下的库目录下。
4、使用:
from efficientnet_pytorch import EfficientNet
model=EfficientNet.from_local_file('efficientnet-b3',<你的模型文件路径>)
之后就可以用该model进行前向推理了。EfficientNet果然不凡,参数量小,速度相当快,网上随意下了几张图片,都能准确推理出来。随后会将代码和模型文件上传到csdn