train_scratch.py

# import debugpy; debugpy.connect(('0.0.0.0', 5681))
import torch
from annotatedtransformer.utils import *
from annotatedtransformer.datasets import Multi30k, Translation2019zh
from annotatedtransformer.vocab import Tokenizer
from annotatedtransformer.transformer import make_model
from torch.utils.data import DataLoader
from torch.optim.lr_scheduler import LambdaLR
import time

def create_dataloaders_translation2019zh(
    vocab_src, vocab_tgt,
    data_dir = 'data/translation2019zh',
    batch_size=12000,
    language_pair=("en", "zh")
):
    def collate_fn(batch):
        return collate_batch(batch, vocab_src, vocab_tgt, vocab_src.pad_idx)

    def get_loader(split):
        assert split in ["train", "valid"]
        dataset = Translation2019zh(data_dir, split, language_pair)
        return DataLoader(
            dataset, batch_size, collate_fn=collate_fn, 
            shuffle=True if split == "train" else False
        )

    train_dataloader = get_loader("train")
    valid_dataloader = get_loader("valid")
    return train_dataloader, valid_dataloader

def create_dataloaders_Multi30k(
    vocab_src, vocab_tgt,
    data_dir = 'data/multi30k',
    batch_size=12000,
    language_pair=("de", "en")
):
    def collate_fn(batch):
        return collate_batch(batch, vocab_src, vocab_tgt, vocab_src.pad_idx)

    def get_loader(split):
        assert split in ["train", "val", "test"]
        dataset = Multi30k(data_dir, split, language_pair)
        return DataLoader(
            dataset, batch_size, collate_fn=collate_fn, 
            shuffle=True if split == "train" else False
        )

    train_dataloader = get_loader("train")
    valid_dataloader = get_loader("val")
    test_dataloader = get_loader("test")
    return train_dataloader, valid_dataloader, test_dataloader

class TrainState:
    """Track number of steps, examples, and tokens processed"""
    step: int = 0  # Steps in the current epoch
    accum_step: int = 0  # Number of gradient accumulation steps
    samples: int = 0  # total # of examples used
    tokens: int = 0  # total # of tokens processed

def run_epoch(
    data_iter,
    model,
    loss_compute,
    optimizer,
    scheduler,
    mode="train",
    accum_iter=1,
    train_state=TrainState(),
):
    """Train a single epoch"""
    start = time.time()
    total_tokens = 0
    total_loss = 0
    tokens = 0
    n_accum = 0
    for i, batch in enumerate(data_iter):
        out = model.forward(
            batch.src, batch.tgt, batch.src_mask, batch.tgt_mask
        )
        loss, loss_node = loss_compute(out, batch.tgt_y, batch.ntokens)
        if mode == "train" or mode == "train+log":
            loss_node.backward()
            train_state.step += 1
            train_state.samples += batch.src.shape[0]
            train_state.tokens += batch.ntokens
            if i % accum_iter == 0:
                optimizer.step()
                optimizer.zero_grad(set_to_none=True)
                n_accum += 1
                train_state.accum_step += 1
            scheduler.step()

        total_loss += loss
        total_tokens += batch.ntokens
        tokens += batch.ntokens
        if i % 40 == 1 and (mode == "train" or mode == "train+log"):
            lr = optimizer.param_groups[0]["lr"]
            elapsed = time.time() - start
            print(
                (
                    "Epoch Step: %6d | Accumulation Step: %3d | Loss: %6.2f "
                    + "| Tokens / Sec: %7.1f | Learning Rate: %6.1e"
                )
                % (i, n_accum, loss / batch.ntokens, tokens / elapsed, lr),
                flush=True,
            )
            start = time.time()
            tokens = 0
        del loss
        del loss_node
    return total_loss / total_tokens, train_state

def train():
    outdir = "output/model/en2de"
    from pathlib import Path
    Path(outdir).mkdir(parents=True, exist_ok=True)

    # config
    batch_size = 256
    d_model = 512
    base_lr = 1.0
    num_epochs = 20
    accum_iter = 10
    file_prefix = f"{outdir}/transformer_"
    lang_pair = ("en", "de")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print("Device:", device)
    
    vocab_src = Tokenizer("output/vocab/de-en/bpe/en_8000.model")
    vocab_tgt = Tokenizer("output/vocab/de-en/bpe/de_8000.model")

    pad_idx = vocab_tgt.pad_idx
    model = make_model(len(vocab_src), len(vocab_tgt), N=6, d_model=d_model)
    model = model.to(device)

    criterion = LabelSmoothing(len(vocab_tgt), pad_idx, 0.1)
    # criterion = nn.CrossEntropyLoss(reduction="sum", ignore_index=pad_idx)

    train_dataloader, valid_dataloader, _ = create_dataloaders_Multi30k(
        vocab_src, vocab_tgt,
        batch_size=batch_size,
        language_pair=lang_pair
    )
    optimizer = torch.optim.Adam(
        model.parameters(), lr=base_lr, betas=(0.9, 0.98), eps=1e-9
    )

    warmup = int ( num_epochs * len(train_dataloader) * 0.2 )
    print("Warmup Steps:", warmup)
    lr_scheduler = LambdaLR(
        optimizer=optimizer,
        lr_lambda=lambda step: rate(
            step, d_model, factor=1, warmup=warmup
        ),
    )

    train_state = TrainState()

    for epoch in range(num_epochs):
        model.train()
        print(f"Epoch {epoch} Training ====", flush=True)
        _, train_state = run_epoch(
            (Batch(b[0].to(device), b[1].to(device), pad_idx) for b in train_dataloader),
            model,
            SimpleLossCompute(model.generator, criterion),
            optimizer,
            lr_scheduler,
            mode="train+log",
            accum_iter=accum_iter,
            train_state=train_state,
        )

        torch.cuda.empty_cache()

        print(f"Epoch {epoch} Validation ====", flush=True)
        model.eval()
        sloss, _ = run_epoch(
            (Batch(b[0].to(device), b[1].to(device), pad_idx) for b in valid_dataloader),
            model,
            SimpleLossCompute(model.generator, criterion),
            DummyOptimizer(),
            DummyScheduler(),
            mode="eval",
        )
        print(sloss)
        torch.cuda.empty_cache()

    file_path = file_prefix + "model.pt"
    torch.save(model.state_dict(), file_path)
    print(f"Model saved to {file_path}")

    
def main():
    train()

if __name__ == "__main__":
    main()