EfficientNet-B0: pizza vs steak classifier

A binary image classifier built with transfer learning, as a beginner learning project.

Results

Test accuracy: 99.40%

Model

Base: EfficientNet-B0 pretrained on ImageNet (torchvision)
New head: Dropout(0.2) -> Linear(1280, 1) -> sigmoid
Output: probability that the image is class index 1 (steak); class index 0 is pizza

Training

Phase 1: frozen backbone, new head only (Adam, lr 0.001, 5 epochs)
Phase 2: last 3 feature blocks unfrozen (Adam, lr 0.0001, 5 epochs)
Loss: BCEWithLogitsLoss, batch size 32
Augmentation: random resized crop, horizontal flip, color jitter
Training images: 1200, validation: 300, test: 500

How to use

from huggingface_hub import hf_hub_download
import sys

repo = "Rhoda07/efficientnet-b0-pizza-steak"
weights = hf_hub_download(repo, "efficientnet_b0_pizza_steak.pth")
config = hf_hub_download(repo, "config.json")
code = hf_hub_download(repo, "model.py")

import importlib.util
spec = importlib.util.spec_from_file_location("model", code)
m = importlib.util.module_from_spec(spec)
spec.loader.exec_module(m)

model = m.load_model(weights)
print(m.predict(model, "photo.jpg", config))

Limitations

The model only knows two classes and will always answer with one of them, even for unrelated photos.

Downloads last month: -