Update file train

model_efficientnet.py

@@ -0,0 +1,15 @@
+import torch.nn as nn
+from torchvision.models import efficientnet_b0, EfficientNet_B0_Weights
+
+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)

train_efficientnet.py

@@ -0,0 +1,78 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+from model_efficientnet import CatDogEfficientNetB0  
+from tqdm import tqdm  # Thêm tqdm để hiển thị tiến trình
+
+# Cấu hình
+BATCH_SIZE = 32
+EPOCHS = 10
+LR = 0.001
+MOMENTUM = 0.9
+WEIGHT_DECAY = 0.0001
+
+# Tiền xử lý dữ liệu
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+train_dataset = datasets.ImageFolder('data/train', transform=transform)
+val_dataset = datasets.ImageFolder('data/val', transform=transform)
+
+train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
+val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)
+
+# Load mô hình EfficientNet từ file model_efficientnet.py
+model = CatDogEfficientNetB0()
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model.parameters(), lr=LR)
+# optimizer = optim.SGD(model.parameters(), lr=LR, momentum=MOMENTUM, weight_decay=WEIGHT_DECAY)
+best_acc = 0.0  # Biến lưu val acc tốt nhất
+
+# Train loop
+for epoch in range(EPOCHS):
+    model.train()
+    running_loss = 0.0
+    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()
+        running_loss += loss.item() * images.size(0)
+        train_bar.set_postfix(loss=loss.item())
+    epoch_loss = running_loss / len(train_loader.dataset)
+    print(f"Epoch {epoch+1}/{EPOCHS}, Loss: {epoch_loss:.4f}")
+
+    # Đánh giá trên tập validation
+    model.eval()
+    correct = 0
+    total = 0
+    with torch.no_grad():
+        for images, labels in val_loader:
+            images, labels = images.to(device), labels.to(device)
+            outputs = model(images)
+            _, preds = torch.max(outputs, 1)
+            correct += (preds == labels).sum().item()
+            total += labels.size(0)
+    acc = correct / total
+    print(f"Validation Accuracy: {acc:.4f}")
+
+    # Lưu checkpoint nếu val acc tốt nhất
+    if acc > best_acc:
+        best_acc = acc
+        torch.save(model.state_dict(), 'efficientnet_best.pth')
+        print(f"==> Đã lưu model tốt nhất với val acc: {best_acc:.4f}")
+
+# Lưu model cuối cùng
+torch.save(model.state_dict(), 'efficientnet_model_final.pth')