experimental code

Author: rasbt

.other/try-accelerate.py

@@ -0,0 +1,191 @@
+import os
+import os.path as op
+import time
+
+from datasets import load_dataset
+import torch
+from torch.utils.data import DataLoader
+import torchmetrics
+from transformers import AutoTokenizer
+from transformers import AutoModelForSequenceClassification
+from watermark import watermark
+from accelerate import Accelerator
+
+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 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}%"
+)
+train_acc.reset(), val_acc.reset()
+
+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"
+accelerator = Accelerator()
+device = accelerator.device
+
+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=4,
+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
+#########################################
+
+model, optimizer, train_loader, val_loader, test_loader = accelerator.prepare(model, optimizer, train_loader, val_loader, test_loader)
+
+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}%")