vanv123/aidd

1import torch
2import torch.nn as nn
3import torch.optim as optim
4from torchvision import datasets, models, transforms
5from torch.utils.data import DataLoader
6import os
7
8# 1. Cấu hình thiết bị (Ưu tiên dùng Card đồ họa)
9device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
10print(f"[INFO] Đang sử dụng thiết bị tính toán: {device}")
11
12# 2. Chuẩn bị và làm phong phú dữ liệu (Data Augmentation)
13# Giúp AI học tốt hơn bằng cách tự động lật ảnh, thay đổi kích thước cho chuẩn
14data_transforms = transforms.Compose([
15    transforms.Resize((224, 224)),
16    transforms.RandomHorizontalFlip(),
17    transforms.ToTensor(),
18    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
19])
20
21# Load dữ liệu từ thư mục
22data_dir = 'dataset/train'
23image_dataset = datasets.ImageFolder(data_dir, data_transforms)
24dataloader = DataLoader(image_dataset, batch_size=16, shuffle=True)
25
26# In ra các nhãn để kiểm tra (thường sẽ là {'fake': 0, 'real': 1})
27print(f"[INFO] Các nhãn đã nhận diện: {image_dataset.class_to_idx}")
28
29# 3. Khởi tạo Mô hình AI (MobileNetV2)
30print("[INFO] Đang tải cấu trúc mạng MobileNetV2...")
31model = models.mobilenet_v2(weights=models.MobileNet_V2_Weights.DEFAULT)
32
33# Thay đổi lớp cuối cùng để chỉ phân loại 2 lớp: Fake và Real
34num_ftrs = model.classifier[1].in_features
35model.classifier[1] = nn.Linear(num_ftrs, 2)
36model = model.to(device) # Đẩy model lên GPU
37
38# 4. Định nghĩa hàm tính sai số (Loss) và bộ tối ưu hóa (Optimizer)
39criterion = nn.CrossEntropyLoss()
40optimizer = optim.Adam(model.parameters(), lr=0.001)
41
42# 5. Vòng lặp Huấn luyện (Training Loop)
43num_epochs = 10 # Cho AI học đi học lại 10 lần qua toàn bộ dữ liệu
44
45print("[INFO] Bắt đầu quá trình huấn luyện...")
46for epoch in range(num_epochs):
47    model.train() # Đặt mô hình ở chế độ huấn luyện
48    running_loss = 0.0
49    corrects = 0
50
51    for inputs, labels in dataloader:
52        inputs = inputs.to(device)
53        labels = labels.to(device)
54
55        optimizer.zero_grad() # Xóa bộ nhớ đệm
56        
57        # Lan truyền tiến (Đưa ảnh qua mạng AI)
58        outputs = model(inputs)
59        _, preds = torch.max(outputs, 1)
60        loss = criterion(outputs, labels)
61
62        # Lan truyền ngược (Cập nhật trọng số để AI thông minh hơn)
63        loss.backward()
64        optimizer.step()
65
66        running_loss += loss.item() * inputs.size(0)
67        corrects += torch.sum(preds == labels.data)
68
69    epoch_loss = running_loss / len(image_dataset)
70    epoch_acc = corrects.double() / len(image_dataset)
71
72    print(f"Vòng lặp (Epoch) {epoch+1}/{num_epochs} | Sai số (Loss): {epoch_loss:.4f} | Độ chính xác (Acc): {epoch_acc:.4f}")
73
74# 6. Lưu lại "bộ não" đã được huấn luyện
75os.makedirs('models', exist_ok=True)
76save_path = 'models/anti_spoofing_model.pth'
77torch.save(model.state_dict(), save_path)
78print(f"[INFO] Huấn luyện hoàn tất! Mô hình đã được lưu tại: {save_path}")