275 lines
8.4 KiB
Python
275 lines
8.4 KiB
Python
"""
|
|
Tahap 4: Training Kata
|
|
|
|
Script ini melatih model TCN untuk mengenali kata dalam bahasa isyarat.
|
|
|
|
Sesuai README:
|
|
- Input: dataset/train/kata/
|
|
- Output: models/kata/model_kata.pt
|
|
- Output tambahan: Accuracy, Loss, Training History
|
|
"""
|
|
|
|
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
|
|
|
|
# Add parent directory to path
|
|
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 KataDataset(Dataset):
|
|
"""Dataset untuk kata."""
|
|
|
|
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:
|
|
sequence = np.load(npy_file)
|
|
self.samples.append(sequence)
|
|
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_kata(data_dir: Path, model_dir: Path, epochs: int = EPOCHS,
|
|
batch_size: int = BATCH_SIZE, learning_rate: float = LEARNING_RATE):
|
|
"""Training model untuk kata."""
|
|
|
|
# Setup device
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
print(f"Using device: {device}")
|
|
|
|
# Create dataset
|
|
print(f"\nLoading dataset from: {data_dir}")
|
|
train_dataset = KataDataset(data_dir / "train" / "kata")
|
|
test_dataset = KataDataset(data_dir / "test" / "kata")
|
|
|
|
print(f"Train samples: {len(train_dataset)}")
|
|
print(f"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())}")
|
|
|
|
# Create data loaders
|
|
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)
|
|
|
|
# Create model
|
|
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()):,}")
|
|
|
|
# Loss and optimizer
|
|
criterion = nn.CrossEntropyLoss()
|
|
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
|
|
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
|
|
optimizer, mode='min', factor=0.5, patience=5
|
|
)
|
|
|
|
# Training loop
|
|
history = {
|
|
"train_loss": [], "train_acc": [],
|
|
"test_loss": [], "test_acc": []
|
|
}
|
|
|
|
best_test_acc = 0
|
|
best_model_state = None
|
|
|
|
print("\n" + "="*60)
|
|
print("Training Kata Model")
|
|
print("="*60)
|
|
|
|
for epoch in range(epochs):
|
|
print(f"\nEpoch {epoch+1}/{epochs}")
|
|
print("-" * 40)
|
|
|
|
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_test_acc:
|
|
best_test_acc = test_acc
|
|
best_model_state = model.state_dict().copy()
|
|
print(f"*** New best model! Test Acc: {best_test_acc:.2f}% ***")
|
|
|
|
# Load best model
|
|
if best_model_state is not None:
|
|
model.load_state_dict(best_model_state)
|
|
|
|
# Save model (sesuai README: models/kata/model_kata.pt)
|
|
model_dir.mkdir(parents=True, exist_ok=True)
|
|
kata_model_dir = model_dir / "kata"
|
|
kata_model_dir.mkdir(exist_ok=True)
|
|
|
|
model_path = kata_model_dir / "model_kata.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()}
|
|
encoder_path = kata_model_dir / "label_encoder.json"
|
|
with open(encoder_path, 'w') as f:
|
|
json.dump({"label_to_idx": train_dataset.label_to_idx, "idx_to_label": idx_to_label}, f, indent=2)
|
|
|
|
# Save history
|
|
history_path = kata_model_dir / "history.json"
|
|
with open(history_path, '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('Kata Model Accuracy')
|
|
plt.xlabel('Epoch')
|
|
plt.ylabel('Accuracy (%)')
|
|
plt.legend()
|
|
plt.grid(True)
|
|
plt.savefig(ACCURACY_DIR / "kata_accuracy.png")
|
|
plt.close()
|
|
print(f"Accuracy plot saved to: {ACCURACY_DIR / 'kata_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('Kata Model Loss')
|
|
plt.xlabel('Epoch')
|
|
plt.ylabel('Loss')
|
|
plt.legend()
|
|
plt.grid(True)
|
|
plt.savefig(LOSS_DIR / "kata_loss.png")
|
|
plt.close()
|
|
print(f"Loss plot saved to: {LOSS_DIR / 'kata_loss.png'}")
|
|
|
|
print("\n" + "="*60)
|
|
print(f"Best Test Accuracy: {best_test_acc:.2f}%")
|
|
print("="*60)
|
|
|
|
return history
|
|
|
|
|
|
def main():
|
|
"""Main entry point."""
|
|
parser = argparse.ArgumentParser(description="Training model kata")
|
|
parser.add_argument("--data-dir", type=str, default=None,
|
|
help="Directory dataset (default: dataset/)")
|
|
parser.add_argument("--model-dir", type=str, default=None,
|
|
help="Directory model (default: models/)")
|
|
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 4: TRAINING KATA")
|
|
print("="*60)
|
|
|
|
train_kata(data_dir, model_dir, args.epochs, args.batch_size, args.lr)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|