import os import torch import argparse from torch import nn from torch.utils.data import DataLoader, random_split from torchmetrics import MeanMetric, MaxMetric import lightning as L from torchmetrics.classification.accuracy import Accuracy from torch.optim.lr_scheduler import LinearLR, CosineAnnealingLR, SequentialLR from tokenizers import Tokenizer from src.config import Config from src.model import TranslateModel from src.dataset import TranslateDataset from lightning.pytorch.loggers import TensorBoardLogger from lightning.pytorch.callbacks import RichProgressBar from lightning.pytorch.callbacks import ModelCheckpoint import argparse def parser_args(): parser = argparse.ArgumentParser(description="Training configuration") parser.add_argument("--encoder_layer", type=int, default=6, help="Number of encoder layers") parser.add_argument("--decoder_layer", type=int, default=6, help="Number of decoder layers") parser.add_argument("--embed_dim", type=int, default=512, help="Embedding dimension size") parser.add_argument("--num_heads", type=int, default=8, help="Number of attention heads") parser.add_argument("--drop_out", type=float, default=0.1, help="Dropout rate") parser.add_argument("--max_len", type=int, default=256, help="Maximum sequence length") parser.add_argument("--vocab_size", type=int, default=30000, help="Vocabulary size") parser.add_argument("--wmt_zh_en_path", type=str, default="data/wmt_zh_en_training_corpus.csv", help="Path to WMT zh-en training corpus") parser.add_argument("--tokenizer_file", type=str, default="checkpoints/tokenizer.json", help="Path to tokenizer file") parser.add_argument("--batch_size", type=int, default=64, help="Batch size for training") parser.add_argument("--compile", action="store_true", help="Enable torch.compile if available") parser.add_argument("--seed", type=int, default=42, help="Random seed") parser.add_argument("--val_ratio", type=float, default=0.1, help="Validation data ratio") parser.add_argument("--num_workers", type=int, default=4, help="Number of data loading workers") parser.add_argument("--pin_memory", action="store_true", help="Use pinned memory in dataloader") parser.add_argument("--tensorboard_dir", type=str, default="log/tensorboard", help="Directory for tensorboard logs") parser.add_argument("--checkpoint_dir", type=str, default="log/checkpoint", help="Directory for saving checkpoints") # -------- optimizer / scheduler 参数 -------- parser.add_argument("--base_lr", type=float, default=3e-4, help="Base learning rate") parser.add_argument("--betas", type=float, nargs=2, default=(0.9, 0.98), help="AdamW betas") parser.add_argument("--eps", type=float, default=1e-9, help="AdamW epsilon") parser.add_argument("--weight_decay", type=float, default=0.1, help="Weight decay") parser.add_argument("--warmup_ratio", type=float, default=0.0005, help="Warmup ratio (fraction of total steps)") parser.add_argument("--start_factor", type=float, default=1e-3, help="Linear warmup start factor (relative LR scale)") parser.add_argument("--end_factor", type=float, default=1.0, help="Linear warmup end factor (relative LR scale)") parser.add_argument("--eta_min", type=float, default=3e-6, help="Minimum LR in cosine annealing") parser.add_argument("--max_epochs", type=int, default=10, help="Number of epochs to train") # -------- dataset cache 参数 -------- parser.add_argument("--data_cache_dir", type=str, default="data/cache.pickle", help="Path to cache file for dataset") parser.add_argument("--use_cache", action="store_true", help="Enable dataset caching with pickle") parser.add_argument("--every_n_train_steps", type=int, default=10000, help="Save checkpoint every N training steps") return parser.parse_args() def merge_args_config(config: Config, args): for k, v in vars(args).items(): setattr(config, k, v) return config class TranslateLitModule(L.LightningModule): def __init__(self, config: Config): super().__init__() tokenizer: Tokenizer = Tokenizer.from_file(config.tokenizer_file) self.pad_id = tokenizer.token_to_id("[PAD]") self.net = TranslateModel(config=config) self.train_loss = MeanMetric() self.train_acc = Accuracy( task="multiclass", num_classes=config.vocab_size, ignore_index=self.pad_id ) self.criterion = nn.CrossEntropyLoss(ignore_index=self.pad_id) self.val_loss = MeanMetric() self.val_acc = Accuracy( task="multiclass", num_classes=config.vocab_size, ignore_index=self.pad_id ) self.test_loss = MeanMetric() self.test_acc = Accuracy( task="multiclass", num_classes=config.vocab_size, ignore_index=self.pad_id ) self.val_acc_best = MaxMetric() self.config = config def forward(self, batch) -> torch.Tensor: pred = self.net.forward( src=batch["src"], tgt=batch["tgt"], src_pad_mask=batch["src_pad_mask"], tgt_pad_mask=batch["tgt_pad_mask"], ) return pred def on_train_start(self) -> None: self.train_loss.reset() self.train_acc.reset() self.val_loss.reset() self.val_acc.reset() self.val_acc_best.reset() self.test_loss.reset() self.test_acc.reset() def model_step(self, batch): logits = self.forward(batch) B, L, C = logits.shape loss = self.criterion(logits.reshape(-1, C), batch["label"].reshape(-1)) preds = torch.argmax(logits, dim=-1) return loss, preds.reshape(-1), batch["label"].reshape(-1) def training_step(self, batch, batch_idx): loss, preds, targets = self.model_step(batch) # update and log metrics self.train_loss(loss) self.train_acc(preds, targets) self.log( "train/loss", self.train_loss, on_step=True, on_epoch=True, prog_bar=True ) self.log( "train/acc", self.train_acc, on_step=True, on_epoch=True, prog_bar=True ) # return loss or backpropagation will fail return loss def on_train_epoch_end(self) -> None: pass def validation_step(self, batch, batch_idx: int) -> None: loss, preds, targets = self.model_step(batch) # update and log metrics self.val_loss(loss) self.val_acc(preds, targets) self.log("val/loss", self.val_loss, on_step=True, on_epoch=True, prog_bar=True) self.log("val/acc", self.val_acc, on_step=True, on_epoch=True, prog_bar=True) def on_validation_epoch_end(self) -> None: acc = self.val_acc.compute() # get current val acc self.val_acc_best(acc) # update best so far val acc self.log( "val/acc_best", self.val_acc_best.compute(), sync_dist=True, prog_bar=True ) def test_step(self, batch, batch_idx: int) -> None: loss, preds, targets = self.model_step(batch) # update and log metrics self.test_loss(loss) self.test_acc(preds, targets) self.log( "test/loss", self.test_loss, on_step=False, on_epoch=True, prog_bar=True ) self.log("test/acc", self.test_acc, on_step=False, on_epoch=True, prog_bar=True) def on_test_epoch_end(self) -> None: """Lightning hook that is called when a test epoch ends.""" pass def setup(self, stage: str) -> None: if self.config.compile and stage == "fit": self.net = torch.compile(self.net) def configure_optimizers(self): optimizer = torch.optim.AdamW( self.parameters(), lr=self.config.base_lr, betas=self.config.betas, eps=self.config.eps, weight_decay=self.config.weight_decay, ) total_steps = self.trainer.estimated_stepping_batches warmup_steps = max(1, int(self.config.warmup_ratio * total_steps)) cosine_steps = max(1, total_steps - warmup_steps) warmup = LinearLR( optimizer, start_factor=self.config.start_factor, end_factor=self.config.end_factor, total_iters=warmup_steps, ) cosine = CosineAnnealingLR( optimizer, T_max=cosine_steps, eta_min=self.config.eta_min, ) scheduler = SequentialLR(optimizer, schedulers=[warmup, cosine], milestones=[warmup_steps]) return { "optimizer": optimizer, "lr_scheduler": { "scheduler": scheduler, "interval": "step", # 每个 step 调整 "frequency": 1 } } def prepare_dataloader(dataset, config: Config, shuffle=True): dataloader = DataLoader( dataset, batch_size=config.batch_size, shuffle=shuffle, num_workers=config.num_workers, pin_memory=config.pin_memory, ) return dataloader def prepare_dataset(config: Config): full_ds = TranslateDataset(config=config) val_ratio = config.val_ratio val_len = int(len(full_ds) * val_ratio) train_len = len(full_ds) - val_len train_ds, val_ds = random_split( full_ds, [train_len, val_len], generator=torch.Generator().manual_seed(config.seed), ) return prepare_dataloader(train_ds, config), prepare_dataloader( val_ds, config, False ) def prepare_callback(config: Config): logger = TensorBoardLogger(save_dir=config.tensorboard_dir, name="runs") rich_progress_bar = RichProgressBar() checkpoint = ModelCheckpoint( dirpath=config.checkpoint_dir, filename="translate-{step:05d}", save_weights_only=True, every_n_train_steps=config.every_n_train_steps, save_top_k=-1, ) return logger, [rich_progress_bar, checkpoint] def main(): args = parser_args() config = merge_args_config(Config(), args) L.seed_everything(config.seed) train_loader, val_loader = prepare_dataset(config) model = TranslateLitModule(config=config) logger, callbacks = prepare_callback(config) trainer = L.Trainer(callbacks=callbacks, logger=logger, max_epochs=config.max_epochs) trainer.fit(model=model, train_dataloaders=train_loader, val_dataloaders=val_loader) if __name__ == "__main__": main()