File size: 6,322 Bytes
2aee49a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 | import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import time
import sys
import pickle
import os
from sklearn.model_selection import train_test_split
from model import make_model
sys.stdout.reconfigure(encoding='utf-8')
DATA_PATH = 'data/processed/data.pkl'
VOCAB_PATH = 'data/processed/vocab.pkl'
SAVE_PATH = 'models/model.pt'
BATCH_SIZE = 64
D_MODEL = 192
N_LAYER = 3
N_HEAD = 8
DROPOUT = 0.3 # Tăng từ 0.2 -> 0.3 để giảm overfitting
EPOCHS = 200
LR = 5e-4
WARMUP_EPOCHS = 10
class EarlyStopping:
def __init__(self, patience=50, delta=0.0):
self.patience = patience
self.delta = delta
self.best_acc = 0.0 # track accuracy (higher = better)
self.counter = 0
self.early_stop = False
def __call__(self, val_acc):
if val_acc > self.best_acc + self.delta:
self.best_acc = val_acc
self.counter = 0
else:
self.counter += 1
if self.counter >= self.patience:
self.early_stop = True
class SeqDataset(Dataset):
def __init__(self, sequences):
self.sequences = sequences
def __len__(self):
return len(self.sequences)
def __getitem__(self, idx):
return torch.tensor(self.sequences[idx])
def collate_fn(batch, pad_idx):
seqs = nn.utils.rnn.pad_sequence(batch, batch_first=True, padding_value=pad_idx)
x = seqs[:, :-1]
y = seqs[:, 1:].clone()
y[y == pad_idx] = -1
return x, y
def train():
device = torch.device("cpu")
if not os.path.exists(DATA_PATH): return
with open(VOCAB_PATH, 'rb') as f: vocab = pickle.load(f)
with open(DATA_PATH, 'rb') as f: data = pickle.load(f)
word2idx, pad_idx, vocab_size = vocab['word2idx'], vocab['word2idx']['<pad>'], len(vocab['word2idx'])
train_seq = data['train']['sequences']
val_seq = data['val']['sequences']
train_dl = DataLoader(SeqDataset(train_seq), batch_size=BATCH_SIZE, shuffle=True, collate_fn=lambda b: collate_fn(b, pad_idx))
val_dl = DataLoader(SeqDataset(val_seq), batch_size=BATCH_SIZE, shuffle=False, collate_fn=lambda b: collate_fn(b, pad_idx))
model = make_model(vocab_size, n_layer=N_LAYER, d_model=D_MODEL, n_head=N_HEAD, dropout=DROPOUT)
model.to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=0.05)
# Warmup then ReduceLROnPlateau: LR tự giảm khi val_acc không tăng
warmup_scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=0.1, end_factor=1.0, total_iters=WARMUP_EPOCHS)
plateau_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=15, min_lr=1e-6, verbose=True)
early_stopping = EarlyStopping(patience=60, delta=0.1) # acc-based, stop if no +0.1% in 60 epochs
best_val_loss = float('inf')
best_val_acc = 0
best_epoch = 0
os.makedirs('models', exist_ok=True)
os.makedirs('logs', exist_ok=True)
with open('logs/progress.txt', 'w', encoding='utf-8') as f:
f.write("=== Training Log ===\n")
for epoch in range(EPOCHS):
model.train()
total_train_loss = total_train_acc = 0
start_time = time.time()
for x, y in train_dl:
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
logits, loss = model(x, y)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) # Gradient Clipping
optimizer.step()
total_train_loss += loss.item()
mask = y != -1
total_train_acc += (logits.argmax(dim=-1) == y)[mask].float().mean().item()
if epoch < WARMUP_EPOCHS:
warmup_scheduler.step()
avg_train_loss, avg_train_acc = total_train_loss / len(train_dl), total_train_acc / len(train_dl) * 100
model.eval()
total_val_loss = total_val_acc = 0
with torch.no_grad():
for x, y in val_dl:
x, y = x.to(device), y.to(device)
logits, loss = model(x, y)
total_val_loss += loss.item()
mask = y != -1
total_val_acc += (logits.argmax(dim=-1) == y)[mask].float().mean().item()
avg_val_loss, avg_val_acc = total_val_loss / len(val_dl), total_val_acc / len(val_dl) * 100
is_overfitting = avg_val_loss > avg_train_loss + 2.5
warn = " ⚠️ OVERFIT" if is_overfitting else ""
msg = f"Epoch {epoch+1:3d} | Train: {avg_train_loss:.4f} ({avg_train_acc:.1f}%) | Val: {avg_val_loss:.4f} ({avg_val_acc:.1f}%) | {time.time()-start_time:.1f}s{warn}"
print(msg)
with open('logs/progress.txt', 'a', encoding='utf-8') as f:
f.write(msg + "\n")
if avg_val_loss < best_val_loss:
best_val_loss = avg_val_loss
best_val_acc = avg_val_acc
best_epoch = epoch + 1
torch.save({'model_state_dict': model.state_dict(), 'v_size': vocab_size, 'd_m': D_MODEL, 'n_l': N_LAYER, 'n_h': N_HEAD, 'dr': DROPOUT}, SAVE_PATH)
if is_overfitting:
stop_msg = f"Stopping training due to Overfitting (Gap > 2.5) at epoch {epoch+1}."
print(stop_msg)
with open('logs/progress.txt', 'a', encoding='utf-8') as f:
f.write(stop_msg + "\n")
break
plateau_scheduler.step(avg_val_acc) # ReduceLROnPlateau theo val_acc
early_stopping(avg_val_acc) # track accuracy
if early_stopping.early_stop:
stop_msg = f"Early stopping at epoch {epoch+1}. (no val acc improvement > 0.1% in {early_stopping.patience} epochs)"
print(stop_msg)
with open('logs/progress.txt', 'a', encoding='utf-8') as f:
f.write(stop_msg + "\n")
break
final_msg = f"\n--- Training Complete ---\nBest Epoch: {best_epoch}\nBest Val Loss: {best_val_loss:.4f}\nBest Val Acc: {best_val_acc:.1f}%\nModel saved to: {SAVE_PATH}"
print(final_msg)
with open('logs/progress.txt', 'a', encoding='utf-8') as f:
f.write(final_msg + "\n")
if __name__ == "__main__":
train() |