217 lines
7.6 KiB
Python
217 lines
7.6 KiB
Python
"""
|
|
Tahap 10: Training Angka
|
|
|
|
Script ini melatih model TCN untuk mengenali angka dalam bahasa isyarat.
|
|
|
|
Sesuai README:
|
|
- Input: dataset/train/angka/
|
|
- Output: models/angka/model_angka.pt
|
|
"""
|
|
|
|
import sys
|
|
import json
|
|
import numpy as np
|
|
from pathlib import Path
|
|
from typing import Tuple
|
|
import argparse
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
from torch.utils.data import DataLoader, Dataset
|
|
from tqdm import tqdm
|
|
import matplotlib.pyplot as plt
|
|
|
|
sys.path.insert(0, str(Path(__file__).parent.parent))
|
|
|
|
from utils.tcn import TCNClassifier
|
|
from utils.config import TRAIN_DIR, MODEL_DIR, FEATURE_SIZE, BATCH_SIZE, EPOCHS, LEARNING_RATE, ACCURACY_DIR, LOSS_DIR
|
|
|
|
|
|
class AngkaDataset(Dataset):
|
|
"""Dataset untuk angka."""
|
|
|
|
def __init__(self, data_dir: Path):
|
|
self.data_dir = data_dir
|
|
self.samples = []
|
|
self.labels = []
|
|
self.label_to_idx = {}
|
|
self._load_data()
|
|
|
|
def _load_data(self):
|
|
if not self.data_dir.exists():
|
|
return
|
|
label_dirs = sorted([d for d in self.data_dir.iterdir() if d.is_dir()])
|
|
self.label_to_idx = {d.name: idx for idx, d in enumerate(label_dirs)}
|
|
for label_dir in label_dirs:
|
|
label_idx = self.label_to_idx[label_dir.name]
|
|
for npy_file in label_dir.glob("*.npy"):
|
|
try:
|
|
self.samples.append(np.load(npy_file))
|
|
self.labels.append(label_idx)
|
|
except Exception as e:
|
|
print(f"Error loading {npy_file}: {e}")
|
|
|
|
def __len__(self):
|
|
return len(self.samples)
|
|
|
|
def __getitem__(self, idx) -> Tuple:
|
|
return torch.FloatTensor(self.samples[idx]), torch.LongTensor([self.labels[idx]]).squeeze()
|
|
|
|
|
|
def train_epoch(model, train_loader, criterion, optimizer, device):
|
|
model.train()
|
|
total_loss, correct, total = 0, 0, 0
|
|
for sequences, labels in tqdm(train_loader, desc="Training"):
|
|
sequences, labels = sequences.to(device), labels.to(device)
|
|
optimizer.zero_grad()
|
|
outputs = model(sequences)
|
|
loss = criterion(outputs, labels)
|
|
loss.backward()
|
|
optimizer.step()
|
|
total_loss += loss.item()
|
|
_, predicted = torch.max(outputs, 1)
|
|
total += labels.size(0)
|
|
correct += (predicted == labels).sum().item()
|
|
return total_loss / len(train_loader), correct / total * 100
|
|
|
|
|
|
def evaluate(model, eval_loader, criterion, device):
|
|
model.eval()
|
|
total_loss, correct, total = 0, 0, 0
|
|
with torch.no_grad():
|
|
for sequences, labels in eval_loader:
|
|
sequences, labels = sequences.to(device), labels.to(device)
|
|
outputs = model(sequences)
|
|
loss = criterion(outputs, labels)
|
|
total_loss += loss.item()
|
|
_, predicted = torch.max(outputs, 1)
|
|
total += labels.size(0)
|
|
correct += (predicted == labels).sum().item()
|
|
return total_loss / len(eval_loader), correct / total * 100
|
|
|
|
|
|
def train_angka(data_dir: Path, model_dir: Path, epochs: int = EPOCHS,
|
|
batch_size: int = BATCH_SIZE, lr: float = LEARNING_RATE):
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
print(f"Using device: {device}")
|
|
|
|
print(f"\nLoading dataset from: {data_dir}")
|
|
train_dataset = AngkaDataset(data_dir / "train" / "angka")
|
|
test_dataset = AngkaDataset(data_dir / "test" / "angka")
|
|
|
|
print(f"Train samples: {len(train_dataset)}, Test samples: {len(test_dataset)}")
|
|
|
|
if len(train_dataset) == 0:
|
|
raise ValueError("No training samples found!")
|
|
|
|
num_classes = len(train_dataset.label_to_idx)
|
|
print(f"Number of classes: {num_classes}")
|
|
print(f"Classes: {list(train_dataset.label_to_idx.keys())}")
|
|
|
|
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
|
|
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)
|
|
|
|
model = TCNClassifier(input_size=FEATURE_SIZE, num_classes=num_classes, num_channels=[64, 128, 256]).to(device)
|
|
|
|
print(f"\nTotal parameters: {sum(p.numel() for p in model.parameters()):,}")
|
|
|
|
criterion = nn.CrossEntropyLoss()
|
|
optimizer = optim.Adam(model.parameters(), lr=lr)
|
|
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5)
|
|
|
|
history = {"train_loss": [], "train_acc": [], "test_loss": [], "test_acc": []}
|
|
best_acc = 0
|
|
|
|
print("\n" + "="*60)
|
|
print("Training Angka Model")
|
|
print("="*60)
|
|
|
|
for epoch in range(epochs):
|
|
print(f"\nEpoch {epoch+1}/{epochs}")
|
|
train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer, device)
|
|
test_loss, test_acc = evaluate(model, test_loader, criterion, device)
|
|
scheduler.step(test_loss)
|
|
history["train_loss"].append(train_loss)
|
|
history["train_acc"].append(train_acc)
|
|
history["test_loss"].append(test_loss)
|
|
history["test_acc"].append(test_acc)
|
|
print(f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}%")
|
|
print(f"Test Loss: {test_loss:.4f} | Test Acc: {test_acc:.2f}%")
|
|
if test_acc > best_acc:
|
|
best_acc = test_acc
|
|
best_state = model.state_dict().copy()
|
|
print(f"*** Best model! Acc: {best_acc:.2f}% ***")
|
|
|
|
if best_state:
|
|
model.load_state_dict(best_state)
|
|
|
|
# Save model
|
|
model_dir.mkdir(parents=True, exist_ok=True)
|
|
angka_model_dir = model_dir / "angka"
|
|
angka_model_dir.mkdir(exist_ok=True)
|
|
model_path = angka_model_dir / "model_angka.pt"
|
|
torch.save(model, model_path)
|
|
print(f"\nModel saved to: {model_path}")
|
|
|
|
# Save label encoder
|
|
idx_to_label = {v: k for k, v in train_dataset.label_to_idx.items()}
|
|
with open(angka_model_dir / "label_encoder.json", 'w') as f:
|
|
json.dump({"label_to_idx": train_dataset.label_to_idx, "idx_to_label": idx_to_label}, f, indent=2)
|
|
|
|
# Save history
|
|
with open(angka_model_dir / "history.json", 'w') as f:
|
|
json.dump(history, f, indent=2)
|
|
|
|
# Plot Accuracy
|
|
plt.figure(figsize=(10, 6))
|
|
plt.plot(history['train_acc'], label='Train Accuracy')
|
|
plt.plot(history['test_acc'], label='Test (Validation) Accuracy')
|
|
plt.title('Angka Model Accuracy')
|
|
plt.xlabel('Epoch')
|
|
plt.ylabel('Accuracy (%)')
|
|
plt.legend()
|
|
plt.grid(True)
|
|
plt.savefig(ACCURACY_DIR / "angka_accuracy.png")
|
|
plt.close()
|
|
print(f"Accuracy plot saved to: {ACCURACY_DIR / 'angka_accuracy.png'}")
|
|
|
|
# Plot Loss
|
|
plt.figure(figsize=(10, 6))
|
|
plt.plot(history['train_loss'], label='Train Loss')
|
|
plt.plot(history['test_loss'], label='Test (Validation) Loss')
|
|
plt.title('Angka Model Loss')
|
|
plt.xlabel('Epoch')
|
|
plt.ylabel('Loss')
|
|
plt.legend()
|
|
plt.grid(True)
|
|
plt.savefig(LOSS_DIR / "angka_loss.png")
|
|
plt.close()
|
|
print(f"Loss plot saved to: {LOSS_DIR / 'angka_loss.png'}")
|
|
|
|
print(f"\nBest Test Accuracy: {best_acc:.2f}%")
|
|
return history
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Training model angka")
|
|
parser.add_argument("--data-dir", type=str, default=None)
|
|
parser.add_argument("--model-dir", type=str, default=None)
|
|
parser.add_argument("--epochs", type=int, default=EPOCHS)
|
|
parser.add_argument("--batch-size", type=int, default=BATCH_SIZE)
|
|
parser.add_argument("--lr", type=float, default=LEARNING_RATE)
|
|
args = parser.parse_args()
|
|
|
|
data_dir = Path(args.data_dir) if args.data_dir else TRAIN_DIR.parent
|
|
model_dir = Path(args.model_dir) if args.model_dir else MODEL_DIR
|
|
|
|
print("="*60)
|
|
print("TAHAP 10: TRAINING ANGKA")
|
|
print("="*60)
|
|
|
|
train_angka(data_dir, model_dir, args.epochs, args.batch_size, args.lr)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|