train.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import logging
from cycler import cycler
from lightning.pytorch.loggers import WandbLogger
import umap
from pytorch_metric_learning import losses, miners, reducers, testers
from pytorch_metric_learning.utils.accuracy_calculator import AccuracyCalculator

import lightning as L, torch, torch.nn as nn, torch.nn.functional as F, torchmetrics
from torchmetrics.classification import Accuracy, F1Score 
from torchmetrics import Metric
from torch.utils.data import DataLoader, Dataset, TensorDataset, random_split
from pytorch_metric_learning import distances, losses, miners, reducers, testers, samplers
from pytorch_metric_learning.utils.accuracy_calculator import AccuracyCalculator
from pytorch_metric_learning.utils.inference import InferenceModel
import pandas as pd, numpy as np, matplotlib.pyplot as plt, os, umap, pyarrow as pa, argparse
from pyarrow.parquet import ParquetFile
from tqdm import tqdm 
from rich.console import Console
from rich.table import Table
import pytorch_metric_learning.utils.logging_presets as LP
import wandb
from torchmetrics.classification import Accuracy, ConfusionMatrix
from hf_olmo import *
from transformers import (AdamW,AutoConfig,
                          AutoModelForCausalLM,AutoTokenizer,
                          get_linear_schedule_with_warmup)
from torch.optim.lr_scheduler import CosineAnnealingLR
from utils import *
from model.config import *
from angular_net import *
from typing import Any, List, Dict
from transformers.tokenization_utils_base import BatchEncoding
from angular_net import *

# helper functions
def get_most_labels(train_labels:list, 
                    test_labels:list
                    ) -> list: 
    '''Returns the largest set of class_labels as a list. This is because sometimes we
    get rare labels only in one of the train/test sets'''
    if len(train_labels.unique()) <= len(test_labels.unique()): 
        class_labels = train_labels.unique().int().tolist()
    elif len(train_labels.unique()) > len(test_labels.unique()):
        class_labels = test_labels.unique().int().tolist()
    return class_labels

def map_class_accuracies(class_labels:list, 
                         accuracies:dict, 
                         label_map:dict
                         ) -> dict:
    '''Unpacks accuracies object into dictionary of accuracies by matching with class_labels and label_map'''
    # TODO: check if this is right
    accuracies_dict = {class_labels[i]: accuracies['precision_at_1'][i] 
                       for i in range(len(class_labels)) 
                       if i in range(len(accuracies['precision_at_1']))}
    return accuracies_dict

def plot_confusion_matrix(confusion_matrix:torch.tensor, 
                          label_map:dict, 
                          epoch:int
                          )-> wandb.Image:
    '''Takes in a confusion matrix of numpy ints and label_map dict
    and returns a wandb image for logging'''
    confusion_matrix = confusion_matrix.detach().cpu().numpy().astype('int')
    plt.figure(figsize=(10, 8))
    im = plt.imshow(confusion_matrix, interpolation='nearest', cmap='turbo')
    plt.title(f'Confusion Matrix for Epoch: {epoch}')
    tick_marks = np.arange(len(label_map))
    plt.xticks(tick_marks, [label_map[k] for k in sorted(label_map.keys())], rotation=45)
    plt.yticks(tick_marks, [label_map[k] for k in sorted(label_map.keys())])
    plt.tight_layout()
    plt.ylabel('True label')
    plt.xlabel('Predicted label')
    for i in range(confusion_matrix.shape[0]):
        for j in range(confusion_matrix.shape[1]):
            plt.text(j, i, format(confusion_matrix[i, j], 'd'),
                    ha="center", va="center",
                    color="white" if confusion_matrix[i, j] > im.norm(confusion_matrix).max() / 2. else "white")
    cbr = plt.colorbar()
    cbr.set_label('Count')
    wandb_image_cm = wandb.Image(plt)
    return wandb_image_cm

def get_classification_accuracies(confusion_matrix:torch.Tensor, 
                                  class_labels:list, 
                                  label_map:dict) -> dict: 
    '''Takes in '''
    accuracy_per_class = (confusion_matrix.diag()/confusion_matrix.sum(1)*100).detach().cpu().numpy().tolist()
    class_accuracies_dict = {class_labels[i]: accuracy_per_class[i] for i in range(len(class_labels))}
    mapped_class_accuracies_dict = {}
    for label_name, label_id in label_map.items():
        if label_id in class_accuracies_dict:
            mapped_class_accuracies_dict[label_name] = class_accuracies_dict[label_id]
    return mapped_class_accuracies_dict

def plot_embeddings(umap_embeddings:np.array, 
                    test_labels:list, 
                    label_map:dict, 
                    epoch:int) -> wandb.Image:
    plt.figure(figsize=(10, 8))
    plt.scatter(umap_embeddings[:, 0], umap_embeddings[:, 1], c=test_labels, cmap='Spectral', s=5)
    cbar = plt.colorbar(boundaries=np.arange(len(np.unique(test_labels))+1)-0.5)
    cbar.set_ticks(np.array(np.arange(len(np.unique(test_labels)))))
    cbar.ax.set_yticklabels([v for k, v in label_map.items() if k in test_labels]) # [k for k, v in label_map.items() if v in test_labels]
    plt.title(f'UMAP Projection of the Embeddings, Epoch: {epoch}', fontsize=24)
    wandb_image = wandb.Image(plt)
    return wandb_image


def visualizer_hook(umapper, umap_embeddings, labels, split_name, keyname, *args):
    logging.info(
        "UMAP plot for the {} split and label set {}".format(split_name, keyname)
    )
    label_set = np.unique(labels)
    num_classes = len(label_set)
    plt.figure(figsize=(20, 15))
    plt.gca().set_prop_cycle(
        cycler(
            "color", [plt.cm.nipy_spectral(i) for i in np.linspace(0, 0.9, num_classes)]
        )
    )
    for i in range(num_classes):
        idx = labels == label_set[i]
        plt.plot(umap_embeddings[idx, 0], umap_embeddings[idx, 1], ".", markersize=1)
    plt.show()

def get_all_embeddings(dataset, model):
    tester = testers.BaseTester()
    return tester.get_all_embeddings(dataset, model)

class Metric_Dataset(L.LightningDataModule): 
    def __init__(self, 
                 data_path: str, 
                 num_samples:int, 
                 train_ratio:float, 
                 val_ratio:float, 
                 batch_size:int
                 ): 
        super(Metric_Dataset, self).__init__()
        self.data_path = data_path 
        self.num_samples = num_samples 
        self.train_ratio = train_ratio
        self.val_ratio = val_ratio
        self.batch_size = batch_size
        self.num_workers = int(os.cpu_count() * 0.7)

    def setup(self, stage=None): 
        pf = ParquetFile(self.data_path) 
        first_rows = next(pf.iter_batches(batch_size = self.num_samples)) 
        self.data = pa.Table.from_batches([first_rows]).to_pandas() 
        self.data["embedding"] = self.data["embedding"].apply(lambda x: torch.tensor(x, dtype=torch.float32))
        self.data["one_label"] = self.data["one_label"].apply(lambda x: torch.tensor(x, dtype=torch.float32))
        self.features_tensor = torch.stack(tuple(self.data["embedding"].values))
        self.target_tensor = torch.stack(tuple(self.data["one_label"].values))
        self.dataset = TensorDataset(self.features_tensor, self.target_tensor)
        # split into train_val_test ratio sizes
        total_size = len(self.data)
        train_size = int(self.train_ratio * total_size)
        validation_size = int(train_size * self.val_ratio)
        test_size = total_size - train_size - validation_size
        self.train_dataset, self.validation_dataset, self.test_dataset = random_split(self.dataset, 
                                                                                      [train_size, 
                                                                                       validation_size, 
                                                                                       test_size])
        
    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.batch_size, 
                          shuffle=True, pin_memory=True, num_workers=self.num_workers, drop_last=True)

    def val_dataloader(self):
        return DataLoader(self.validation_dataset, batch_size=self.batch_size, 
                          shuffle=False, pin_memory=True, num_workers=self.num_workers, drop_last=True)

    def test_dataloader(self):
        return DataLoader(self.test_dataset, batch_size=self.batch_size, 
                          shuffle=False, pin_memory=True, num_workers=self.num_workers, drop_last=True)
    
    def input_size(self):
        return self.features_tensor.shape[1]
    
    def num_classes(self): 
        return len(torch.unique(self.target_tensor).tolist())
    
    def return_labels(self):
        return (torch.unique(self.target_tensor)).long().tolist()


if __name__=="__main__": 

    parser = argparse.ArgumentParser(description="pytorch run")
    # training hyperparams
    parser.add_argument('--data_path', default='/lambda_stor/homes/bhsu/gb_2024/my_gb_files/metric_learn/data/arxiv_emb_processed_multi.parquet', 
                        type=str)
    parser.add_argument('--num_epochs', default=15, type=int)
    parser.add_argument('--num_samples', default=10000, type=int)
    parser.add_argument('--validation_ratio', default=0.15, type=float)
    parser.add_argument('--train_ratio', default=0.7, type=float)
    parser.add_argument('--batch_size', default=512, type=int)
    parser.add_argument('--config_path', default = '/homes/bhsu/gb_2024/my_gb_files/metric_learn/metric_learning/training/angular_linear.yaml', type=str)
    parser.add_argument('--lr', default=1e-3, type=float)
    parser.add_argument('--num_devices', default=1, type=int)
    parser.add_argument('--log_offline', default=True, type=bool)
    args = parser.parse_args()

    # model settings 
    metric_cfg = MetricConfig.read_yaml(args.config_path, cfg_type='metric')
    class_cfg = ClassConfig.read_yaml(args.config_path, cfg_type='classifier')
    train_cfg = TrainConfig.read_yaml(args.config_path, cfg_type='training')

    datamodule = Metric_Dataset(data_path=args.data_path, 
                         num_samples=args.num_samples, 
                         train_ratio=0.7,
                         val_ratio=0.15, 
                         batch_size=args.batch_size)

    datamodule.setup()
    train_loader = datamodule.train_dataloader()
    val_loader = datamodule.val_dataloader()
    test_loader = datamodule.test_dataloader()

    # mining + loss functions
    reducer = reducers.ThresholdReducer(low=0)
    metric_loss_fn = losses.ArcFaceLoss(num_classes=datamodule.num_classes(), 
                                                  embedding_size=metric_cfg.output_size, 
                                                  margin=28.6, scale=64)
    mining_fn = miners.AngularMiner(angle=20)
    class_loss_fn = nn.CrossEntropyLoss 
    # accuracy, plots, and logging
    accuracy_calculator = AccuracyCalculator(include=("precision_at_1",
                                                "mean_average_precision"), 
                                                k=1, 
                                                return_per_class=True)
    label_map = dict(zip(datamodule.data["one_label"].apply(int), datamodule.data["one_cat"]))
    umapper = umap.UMAP()
    wandb_logger = WandbLogger(project="metric_lightning", offline=args.log_offline)

    metric_classifier = Metric_Classifier(model_cfg=metric_cfg,
                                   train_cfg=train_cfg,
                                   class_cfg=class_cfg, 
                                   mining_fn=mining_fn,
                                   metric_loss_fn=metric_loss_fn,
                                   class_loss_fn=class_loss_fn, 
                                   train_dataset=datamodule.train_dataset,
                                   val_dataset=datamodule.validation_dataset,
                                   accuracy_calculator=accuracy_calculator,
                                   label_map=label_map, 
                                   umap_embed=umapper
                                   )
    
    trainer = L.Trainer(max_epochs=args.num_epochs, devices=[1], logger=wandb_logger)
    trainer.fit(model=metric_classifier, train_dataloaders=train_loader, val_dataloaders=val_loader)
    
    print("Done")