wzx952/cat-dog-classifier

Cat Dog Classifier

这是一个简单的猫狗分类模型，使用ResNet18架构训练。

模型信息

• 架构: ResNet18 (预训练)
• 输入: 224x224 RGB图像
• 输出: 猫/狗二分类
• 训练框架: PyTorch
• 训练环境: GitHub Actions

使用方法

方法1：直接加载PyTorch模型

import torch
from torchvision import transforms
from PIL import Image

# 加载模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = torch.load("best_model.pth", map_location=device)
model.eval()

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

image = Image.open("your_image.jpg")
input_tensor = transform(image).unsqueeze(0).to(device)

with torch.no_grad():
    output = model(input_tensor)
    probabilities = torch.nn.functional.softmax(output, dim=1)
    predicted_class = "cat" if probabilities[0][0] > 0.5 else "dog"
    confidence = max(probabilities[0][0].item(), probabilities[0][1].item()) * 100

print(f"预测结果: {predicted_class}")
print(f"置信度: {confidence:.2f}%")

方法2：使用Hugging Face Hub

from huggingface_hub import hf_hub_download
import torch

# 下载模型
model_path = hf_hub_download(repo_id="wzx952/cat-dog-classifier", filename="best_model.pth")
model = torch.load(model_path)
model.eval()

训练数据

• 使用标准的猫狗图像数据集进行训练
• 训练集：10张猫图片 + 10张狗图片（演示用）
• 验证集：5张猫图片 + 5张狗图片（演示用）

训练参数

• 优化器: Adam (lr=0.001)
• 损失函数: 交叉熵损失
• 训练轮数: 5
• 批次大小: 32

项目信息

• GitHub仓库: https://github.com/Rhythmzd/ml-project-huggingface
• CI/CD: GitHub Actions自动训练和部署
• Hugging Face: 自动上传模型和文档

演示说明

⚠️ 注意: 这是一个演示项目，使用随机生成的虚拟数据进行训练。实际使用时需要替换为真实的猫狗图像数据集。

许可证

MIT License