CatDogEfficientNetB0 - Phân loại ảnh mèo và chó

Mô tả dự án

Dự án này sử dụng mô hình EfficientNet-B0 để phân loại ảnh mèo và chó. Mô hình được tùy chỉnh để phù hợp với bài toán phân loại 2 lớp (mèo và chó). Dữ liệu được xử lý và huấn luyện thông qua file train_efficientnet.py, và mô hình được định nghĩa trong file model_efficientnet.py.

---

Cấu trúc dự án

• model_efficientnet.py: Định nghĩa mô hình CatDogEfficientNetB0, sử dụng EfficientNet-B0 với lớp cuối được tùy chỉnh để phân loại 2 lớp.
• train_efficientnet.py: File huấn luyện mô hình, bao gồm:
  • Tiền xử lý dữ liệu.
  • Huấn luyện mô hình với thanh tiến trình hiển thị % hoàn thành.
  • Lưu checkpoint khi đạt độ chính xác cao nhất trên tập validation.
• efficientnet_best.pth: Checkpoint của mô hình với độ chính xác cao nhất trên tập validation.
• efficientnet_model_final.pth: Mô hình cuối cùng sau khi hoàn thành tất cả các epoch.

---

Mô hình CatDogEfficientNetB0

Mô hình được định nghĩa trong file model_efficientnet.py:

• Sử dụng EfficientNet-B0 với trọng số ImageNet (EfficientNet_B0_Weights.DEFAULT).
• Đóng băng các trọng số của mô hình gốc (param.requires_grad = False).
• Thay thế lớp phân loại cuối cùng bằng một lớp nn.Linear với 2 đầu ra (mèo và chó).

Code mẫu:

class CatDogEfficientNetB0(nn.Module):
    def __init__(self):
        super().__init__()
        weights = EfficientNet_B0_Weights.DEFAULT
        self.base = efficientnet_b0(weights=weights)
        for param in self.base.parameters():
            param.requires_grad = False
        in_features = self.base.classifier[1].in_features
        self.base.classifier[1] = nn.Linear(in_features, 2)

    def forward(self, x):
        return self.base(x)

---

Huấn luyện mô hình

File train_efficientnet.py thực hiện các bước sau:

1. Tiền xử lý dữ liệu:
   • Resize ảnh về kích thước 224x224.
   • Normalize ảnh theo chuẩn ImageNet.
2. Cấu hình huấn luyện:
   • Sử dụng Adam làm optimizer.
   • Sử dụng CrossEntropyLoss làm hàm mất mát.
   • Huấn luyện trong 10 epoch với batch size 32.
3. Hiển thị tiến trình:
   • Sử dụng thư viện tqdm để hiển thị tiến trình huấn luyện theo từng batch.
4. Lưu checkpoint:
   • Lưu mô hình (efficientnet_best.pth) khi đạt độ chính xác cao nhất trên tập validation.
   • Lưu mô hình cuối cùng (efficientnet_model_final.pth) sau khi hoàn thành tất cả các epoch.

Code mẫu:

for epoch in range(EPOCHS):
    model.train()
    train_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{EPOCHS}", unit="batch")
    for images, labels in train_bar:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        train_bar.set_postfix(loss=loss.item())

---

Cách sử dụng

1. Chuẩn bị dữ liệu:
   • Đặt dữ liệu vào thư mục data/train và data/val theo cấu trúc:

     data/
     ├── train/
     │   ├── cat/
     │   └── dog/
     ├── val/
     │   ├── cat/
     │   └── dog/
     

2. Huấn luyện mô hình:
   • Chạy file train_efficientnet.py:

     python train_efficientnet.py
     

3. Dự đoán:
   • Sử dụng mô hình đã lưu (efficientnet_best.pth) để dự đoán ảnh mới.

---

Yêu cầu

• Python >= 3.8
• Thư viện:
  • torch
  • torchvision
  • tqdm
  • Pillow

---

Kết quả

• Mô hình đạt độ chính xác cao nhất trên tập validation được lưu trong file efficientnet_best.pth.
• Mô hình cuối cùng sau khi hoàn thành tất cả các epoch được lưu trong file efficientnet_model_final.pth.

---

Ghi chú

• Đảm bảo GPU được bật để tăng tốc quá trình huấn luyện.
• Kiểm tra dữ liệu đầu vào để tránh ảnh lỗi hoặc không hợp lệ.