efficientnet-pytorch库使用

pytorch实现的Google最新的imagenet分类模型EfficientNet。

项目地址：https://github.com/lukemelas/EfficientNet-PyTorch

安装

# install
pip install efficientnet_pytorch
# update
pip install --upgrade efficientnet-pytorch

使用

from efficientnet_pytorch import EfficientNet

# 加载预训练模型
model = EfficientNet.from_pretrained('efficientnet-b7', num_classes=54)

# 不使用预训练模型
model = EfficientNet.from_name('efficientnet-b7', override_params={'num_classes': 54})

模型准确度

各个模型在ImageNet数据集上的准确度如下表：

各模型参数

""" Map EfficientNet model name to parameter coefficients. """
    params_dict = {
        # Coefficients:   width,depth,res,dropout
        'efficientnet-b0': (1.0, 1.0, 224, 0.2),
        'efficientnet-b1': (1.0, 1.1, 240, 0.2),
        'efficientnet-b2': (1.1, 1.2, 260, 0.3),
        'efficientnet-b3': (1.2, 1.4, 300, 0.3),
        'efficientnet-b4': (1.4, 1.8, 380, 0.4),
        'efficientnet-b5': (1.6, 2.2, 456, 0.4),
        'efficientnet-b6': (1.8, 2.6, 528, 0.5),
        'efficientnet-b7': (2.0, 3.1, 600, 0.5),
    }

示例：分类

img.jpg和labels_map.txt是测试图像和标签映射文件

import json
from PIL import Image
import torch
from torchvision import transforms

from efficientnet_pytorch import EfficientNet
model = EfficientNet.from_pretrained('efficientnet-b0')

# Preprocess image
tfms = transforms.Compose([transforms.Resize(224), transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),])
img = tfms(Image.open('img.jpg')).unsqueeze(0)
print(img.shape) # torch.Size([1, 3, 224, 224])

# Load ImageNet class names
labels_map = json.load(open('labels_map.txt'))
labels_map = [labels_map[str(i)] for i in range(1000)]

# Classify
model.eval()
with torch.no_grad():
    outputs = model(img)

# Print predictions
print('-----')
for idx in torch.topk(outputs, k=5).indices.squeeze(0).tolist():
    prob = torch.softmax(outputs, dim=1)[0, idx].item()
    print('{label:<75} ({p:.2f}%)'.format(label=labels_map[idx], p=prob*100))

示例：特征提取

from efficientnet_pytorch import EfficientNet
model = EfficientNet.from_pretrained('efficientnet-b0')

# ... image preprocessing as in the classification example ...
print(img.shape) # torch.Size([1, 3, 224, 224])

features = model.extract_features(img)
print(features.shape) # torch.Size([1, 1280, 7, 7])

最后

使用预训练模型的时候，不一定要使得图像的分辨率和模型参数中的要求一致，因为网络结构中有个AdaptiveAvgPool2d结构。