quick sync

Author: rasbt

.gitignore

@@ -1,3 +1,4 @@
+.DS_Store
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

1_pytorch-distilbert.py

@@ -0,0 +1,212 @@
+import os
+import os.path as op
+import time
+
+from datasets import load_dataset
+import matplotlib.pyplot as plt
+import pandas as pd
+import torch
+from torch.utils.data import DataLoader
+import torchmetrics
+from transformers import AutoTokenizer
+from transformers import AutoModelForSequenceClassification
+from watermark import watermark
+
+from local_dataset_utilities import (
+    download_dataset,
+    load_dataset_into_to_dataframe,
+    partition_dataset,
+)
+from local_dataset_utilities import IMDBDataset
+
+
+def tokenize_text(batch):
+    return tokenizer(batch["text"], truncation=True, padding=True)
+
+
+def plot_logs(log_dir):
+    metrics = pd.read_csv(op.join(log_dir, "metrics.csv"))
+
+    aggreg_metrics = []
+    agg_col = "epoch"
+    for i, dfg in metrics.groupby(agg_col):
+        agg = dict(dfg.mean())
+        agg[agg_col] = i
+        aggreg_metrics.append(agg)
+
+    df_metrics = pd.DataFrame(aggreg_metrics)
+    df_metrics[["train_loss", "val_loss"]].plot(
+        grid=True, legend=True, xlabel="Epoch", ylabel="Loss"
+    )
+    plt.savefig(op.join(log_dir, "loss.pdf"))
+
+    df_metrics[["train_acc", "val_acc"]].plot(
+        grid=True, legend=True, xlabel="Epoch", ylabel="Accuracy"
+    )
+    plt.savefig(op.join(log_dir, "acc.pdf"))
+
+
+def train(num_epochs, model, optimizer, train_loader, val_loader, device):
+    for epoch in range(num_epochs):
+        train_acc = torchmetrics.Accuracy(task="multiclass", num_classes=2).to(device)
+
+        for batch_idx, batch in enumerate(train_loader):
+            model.train()
+            for s in ["input_ids", "attention_mask", "label"]:
+                batch[s] = batch[s].to(device)
+
+            ### FORWARD AND BACK PROP
+            outputs = model(
+                batch["input_ids"],
+                attention_mask=batch["attention_mask"],
+                labels=batch["label"],
+            )
+            optimizer.zero_grad()
+            outputs["loss"].backward()
+
+            ### UPDATE MODEL PARAMETERS
+            optimizer.step()
+
+            ### LOGGING
+            if not batch_idx % 300:
+                print(
+                    f"Epoch: {epoch+1:04d}/{num_epochs:04d} | Batch {batch_idx:04d}/{len(train_loader):04d} | Loss: {outputs['loss']:.4f}"
+                )
+
+            model.eval()
+            with torch.no_grad():
+                predicted_labels = torch.argmax(outputs["logits"], 1)
+                train_acc.update(predicted_labels, batch["label"])
+
+        ### MORE LOGGING
+        with torch.no_grad():
+            model.eval()
+            val_acc = torchmetrics.Accuracy(task="multiclass", num_classes=2).to(device)
+            for batch in val_loader:
+                for s in ["input_ids", "attention_mask", "label"]:
+                    batch[s] = batch[s].to(device)
+                outputs = model(
+                    batch["input_ids"],
+                    attention_mask=batch["attention_mask"],
+                    labels=batch["label"],
+                )
+                predicted_labels = torch.argmax(outputs["logits"], 1)
+                val_acc.update(predicted_labels, batch["label"])
+
+            print(
+                f"Epoch: {epoch+1:04d}/{num_epochs:04d} | Train acc.: {train_acc.compute()*100:.2f}% | Val acc.: {val_acc.compute()*100:.2f}%"
+            )
+
+
+if __name__ == "__main__":
+    print(watermark(packages="torch,lightning,transformers", python=True))
+    print("Torch CUDA available?", torch.cuda.is_available())
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+    torch.manual_seed(123)
+
+    ##########################
+    ### 1 Loading the Dataset
+    ##########################
+    download_dataset()
+    df = load_dataset_into_to_dataframe()
+    if not (op.exists("train.csv") and op.exists("val.csv") and op.exists("test.csv")):
+        partition_dataset(df)
+
+    imdb_dataset = load_dataset(
+        "csv",
+        data_files={
+            "train": "train.csv",
+            "validation": "val.csv",
+            "test": "test.csv",
+        },
+    )
+
+    #########################################
+    ### 2 Tokenization and Numericalization
+    #########################################
+
+    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
+    print("Tokenizer input max length:", tokenizer.model_max_length, flush=True)
+    print("Tokenizer vocabulary size:", tokenizer.vocab_size, flush=True)
+
+    print("Tokenizing ...", flush=True)
+    imdb_tokenized = imdb_dataset.map(tokenize_text, batched=True, batch_size=None)
+    del imdb_dataset
+    imdb_tokenized.set_format("torch", columns=["input_ids", "attention_mask", "label"])
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+    #########################################
+    ### 3 Set Up DataLoaders
+    #########################################
+
+    train_dataset = IMDBDataset(imdb_tokenized, partition_key="train")
+    val_dataset = IMDBDataset(imdb_tokenized, partition_key="validation")
+    test_dataset = IMDBDataset(imdb_tokenized, partition_key="test")
+
+    train_loader = DataLoader(
+        dataset=train_dataset,
+        batch_size=12,
+        shuffle=True,
+        num_workers=4,
+        drop_last=True,
+    )
+
+    val_loader = DataLoader(
+        dataset=val_dataset,
+        batch_size=12,
+        num_workers=2,
+        drop_last=True,
+    )
+
+    test_loader = DataLoader(
+        dataset=test_dataset,
+        batch_size=12,
+        num_workers=2,
+        drop_last=True,
+    )
+
+    #########################################
+    ### 4 Initializing the Model
+    #########################################
+
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "distilbert-base-uncased", num_labels=2
+    )
+
+    model.to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)
+
+    #########################################
+    ### 5 Finetuning
+    #########################################
+
+    start = time.time()
+    train(
+        num_epochs=3,
+        model=model,
+        optimizer=optimizer,
+        train_loader=train_loader,
+        val_loader=val_loader,
+        device=device,
+    )
+
+    end = time.time()
+    elapsed = end - start
+    print(f"Time elapsed {elapsed/60:.2f} min")
+
+    with torch.no_grad():
+        model.eval()
+        test_acc = torchmetrics.Accuracy(task="multiclass", num_classes=2).to(device)
+        for batch in test_loader:
+            for s in ["input_ids", "attention_mask", "label"]:
+                batch[s] = batch[s].to(device)
+            outputs = model(
+                batch["input_ids"],
+                attention_mask=batch["attention_mask"],
+                labels=batch["label"],
+            )
+            predicted_labels = torch.argmax(outputs["logits"], 1)
+            test_acc.update(predicted_labels, batch["label"])
+
+    print(f"Test accuracy {test_acc.compute()*100:.2f}%")