|
|
--- |
|
|
license: mit |
|
|
metrics: |
|
|
- accuracy |
|
|
pipeline_tag: image-classification |
|
|
tags: |
|
|
- images |
|
|
--- |
|
|
# 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: |
|
|
```python |
|
|
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: |
|
|
```python |
|
|
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`: |
|
|
```bash |
|
|
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ệ. |
