README.md

metadata

license: mit
tags:
  - image-classification
  - pytorch
  - resnet18
  - transfer-learning
  - computer-vision
pipeline_tag: image-classification

Cat vs Dog Classifier (ResNet18 Transfer Learning)

This model is a fine-tuned ResNet18 for binary image classification: cat vs dog.

Model Details

• Base model: ResNet18 (pretrained on ImageNet)
• Framework: PyTorch
• Task: Binary image classification (cat, dog)
• Input size: 128x128 RGB images
• Training method:
  1. Feature extraction — froze all layers except the final fully connected layer, trained for 5 epochs.
  2. Fine-tuning — unfroze layer4 of ResNet18 and trained further with a lower learning rate for 5 epochs.

Dataset

Trained on a small subset of CIFAR-10 (cat and dog classes only), with 100 images per class for training and 50 images per class for testing.

How to Use

import torch
import torch.nn as nn
from torchvision import models, transforms
from huggingface_hub import hf_hub_download
from PIL import Image

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class_names = {0: "cat", 1: "dog"}

transform = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                          std=[0.229, 0.224, 0.225]),
])

weights_path = hf_hub_download(repo_id="billahaiml/cat-dog-resnet18", filename="cat_dog_resnet18.pth")

model = models.resnet18(weights=None)
model.fc = nn.Linear(model.fc.in_features, 2)
model.load_state_dict(torch.load(weights_path, map_location=device))
model = model.to(device)
model.eval()

img = Image.open("your_image.jpg").convert("RGB")
img_tensor = transform(img).unsqueeze(0).to(device)

with torch.no_grad():
    output = model(img_tensor)
    probs = torch.softmax(output, dim=1)[0]
    pred = class_names[torch.argmax(probs).item()]

print(f"Prediction: {pred}")

Limitations

This model was trained on a very small dataset (100 images per class), so it is intended for educational/demo purposes rather than production use.