cat-dog-classifier

Cat vs Dog Classifier

A PyTorch image classifier that distinguishes cats from dogs, trained on the Kaggle Cat and Dog dataset.

Models Included

Usage

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

# Load model
model = models.resnet18()
model.fc = nn.Linear(model.fc.in_features, 2)
model.load_state_dict(torch.load("transfer_learning_layer4_fc.pth", weights_only=True))
model.eval()

# Preprocess
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
])

image = Image.open("cat.jpg")
tensor = transform(image).unsqueeze(0)

# Predict
with torch.no_grad():
    logits = model(tensor)
    probs = torch.softmax(logits, dim=1)
    predicted = torch.argmax(probs, dim=1).item()

classes = ["cat", "dog"]
print(f"Prediction: {classes[predicted]} ({probs[0][predicted]:.1%})")

Training Details

• Dataset: Kaggle Cat and Dog (8,007 train / 2,025 test)
• Input: 224x224 RGB, normalized to [-1, 1]
• Optimizer: Adam (lr=0.001)
• Loss: CrossEntropyLoss
• Epochs: MLP(3), CNN(5), Transfer Learning(3)

Results