hackathonf23-Creativity_Underflow/train.py

# Standard Library Imports
import argparse
import os
import time

import matplotlib.pyplot as plt
import numpy as np
# External Library Imports
import torch
import torch.nn.functional as F
import tqdm
from PIL import Image
from einops import rearrange
from torch.utils.data import DataLoader
from transformers import AutoProcessor, GroupViTVisionModel

import settings
import wandb
# Local Module Imports
from modelingEntityGrounder import FineGrainedSimilarityLoss, EntityGrounder, ImageProj
from segment_merger import SegmentMerger
from utils.embeddingDataset import EmbeddingDataset
from utils.region_grouper import get_segments

# Constants and Configuration
torch.manual_seed(0)
BATCH_SIZE = 32
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
argparser = argparse.ArgumentParser()
argparser.add_argument("--debug", action="store_true")
argparser.add_argument("--wandb", action="store_true")
DEBUG = argparser.parse_args().debug
USE_LOGGING = argparser.parse_args().wandb


class Trainer:
    def __init__(self, config):
        proj_dim = 256
        self.text_model: EntityGrounder = EntityGrounder(input_dim_text=256, lstm_hidden_dim=256, intermediate_dim=512, output_dim=proj_dim, hidden_dim=512, num_hidden=1).to(
            device)
        self.image_model: ImageProj = ImageProj(196, proj_dim, 512, 1).to(device)

        self.text_optimizer: torch.optim.AdamW = torch.optim.AdamW(self.text_model.parameters(), lr=0.0001, weight_decay=0.05)
        self.image_optimizer: torch.optim.AdamW = torch.optim.AdamW(self.image_model.parameters(), lr=0.0001, weight_decay=0.05)

        self.entity_grounder_loss: FineGrainedSimilarityLoss = FineGrainedSimilarityLoss(-0.5)
        self.segment_merger_loss: SegmentMerger = SegmentMerger(1, 64)

        self.config = config

        if DEBUG:
            self.train_data = EmbeddingDataset("val_data/val_data_images.pk", "val_data/data_text.pk", False, False)
        else:
            self.train_data = EmbeddingDataset("train_data/data_images.pk", "train_data/data_text.pk", True, False)

        self.validation_data = EmbeddingDataset("val_data/val_data_images.pk", "val_data/data_text.pk", False, True)

        self.train_data_loader: DataLoader = DataLoader(self.train_data, batch_size=BATCH_SIZE, shuffle=True)
        self.validation_data_loader: DataLoader = DataLoader(self.validation_data, batch_size=BATCH_SIZE, shuffle=True)

        self.image_processor = AutoProcessor.from_pretrained("nvidia/groupvit-gcc-yfcc")
        self.region_proposer = GroupViTVisionModel.from_pretrained("nvidia/groupvit-gcc-yfcc")

        ### metrics
        self.train_steps: int = 0
        self.epoch: int = 0

    def load_models(self, *, text_model_path=None, image_model_path=None):
        if text_model_path is not None:
            self.text_model.load_state_dict(torch.load(text_model_path))
        if image_model_path is not None:
            self.image_model.load_state_dict(torch.load(image_model_path))

    def forward(self, text_segments, image_segments):
        text_proj = self.text_model(text_segments)
        image_proj = self.image_model(image_segments)
        return image_proj, text_proj

    def backward(self, loss):
        self.text_optimizer.zero_grad()
        self.image_optimizer.zero_grad()
        loss.backward()
        self.text_optimizer.step()
        self.image_optimizer.step()

    def setup(self):
        if USE_LOGGING:
            wandb.init(project=self.config.TrainSettings.Logging.project_name)
            if not os.path.exists(f"{self.config.TrainSettings.Logging.model_dir}"):
                os.mkdir(self.config.TrainSettings.Logging.model_dir)
            os.mkdir(f"{self.config.TrainSettings.Logging.model_dir}/{wandb.run.name}")

    def save(self):
        if USE_LOGGING:
            torch.save(self.text_model.state_dict(), f"{self.config.TrainSettings.Logging.model_dir}/{wandb.run.name}/model_{self.epoch}-{int(time.time())}.pt")
            torch.save(self.image_model.state_dict(), f"{self.config.TrainSettings.Logging.model_dir}/{wandb.run.name}/model_image_proj_{self.epoch}-{int(time.time())}.pt")

    @staticmethod
    def log(args):
        if USE_LOGGING:
            wandb.log(args)


def verifyEntityGrounder(trainer, num_val=10):
    """
    Verification function that processes image and text segments, generates a wandb table, 
    and logs results for visualization.
    
    Args:
    - val_dataloader (DataLoader): The validation data loader.
    - model (torch.nn.Module): The main model for processing.
    - model_image (torch.nn.Module): The image processing model.
    - processor (AutoProcessor): Processor for the GroupViT model.
    - num_val (int): Number of validation instances to process.
    
    Returns:
    None
    """

    index = 0
    with torch.no_grad():
        for image_segs, text_segs, images, strings in trainer.validation_data_loader:
            if image_segs.size(0) != BATCH_SIZE:
                continue
            b = image_segs.size(0)
            text_segs = text_segs.to(device)
            image_segs = image_segs.to(device)
            mask = (text_segs != 0)
            text_segs = rearrange(text_segs, 'b c t d -> (b c) t d')
            image_proj, txt = trainer.forward(text_segs, image_segs)

            txt = rearrange(txt, '(b c) d -> b c d', b=b)

            table = wandb.Table(columns=["image", "heat_map", "caption", "cosine sim p"])
            sims = F.cosine_similarity(image_proj.unsqueeze(2), txt.unsqueeze(1), dim=-1).detach()  # b, seg, text
            mask = rearrange(mask, 'b c t d -> b c (t d)')
            mask = mask.any(dim=-1)  # Reduce the last dimension by checking if there's any non-zero value
            sims = sims * mask.unsqueeze(1).float()
            sims = sims.cpu()
            for i in range(16):
                pic = Image.open(images[i])
                # Get segmentations from attention outputs
                segs = get_segments(pic, trainer.image_processor, trainer.region_proposer).detach().cpu().numpy()

                # Construct tensor to hold image with segments colored by cosSim to the first noun
                color_seg = np.zeros((segs.shape[0], segs.shape[1], 3), dtype=np.uint8)  # height, width, 3

                # Go through each segment, get the cosSim with the first noun, color with colormap
                for seg in range(64):
                    viridis_color = plt.get_cmap("Blues")(sims[i, seg, 0])[0:3]
                    viridis_color = np.array(viridis_color) * 255
                    color_seg[segs == seg, :] = viridis_color

                # create plot of cosSimP (i b/c index into segment)
                fig, ax = plt.subplots()
                cos_p_fig = ax.matshow(sims[i].numpy(), cmap="viridis", vmin=-1, vmax=1)
                fig.colorbar(cos_p_fig)
                fig.savefig("cossimp.png")
                cossimpimg = Image.open("cossimp.png")
                cossimpimg = wandb.Image(cossimpimg, caption="cos sim p")

                plt.close(fig)

                img = wandb.Image(pic, caption=strings[0][i])
                pic = np.array(pic)
                img_to_hmap_scale = 0.1
                if len(pic.shape) == 2:
                    pic = np.expand_dims(pic, axis=2)
                heat_map = np.array(pic * img_to_hmap_scale + color_seg * (1 - img_to_hmap_scale), dtype=np.uint8)
                heat_map = wandb.Image(heat_map)
                table.add_data(img, heat_map, strings[0][i], cossimpimg)
            index += 1
            if index >= num_val:
                break

        trainer.log({"entityGrounderExamples": table})


def trainEntityGrounder(trainer):
    train_loss = 0.0
    pbar = tqdm.tqdm(trainer.train_data_loader, total=len(trainer.train_data_loader))
    index = 0
    for image_segs, text_segs, images, nouns in trainer.train_data_loader:
        if image_segs.size(0) != BATCH_SIZE:
            continue
        image_segs = image_segs.detach().to(device)
        text_segs = text_segs.detach().to(device)

        b = image_segs.size(0)
        mask = (text_segs != 0)
        mask = rearrange(mask, 'b c t d -> b c (t d)')
        mask = mask.any(dim=-1)  # Reduce the last dimension by checking if there's any non-zero value

        text_segs = rearrange(text_segs, 'b c t d -> (b c) t d')
        image_proj, txt_proj = trainer.forward(text_segs, image_segs)
        txt_proj = rearrange(txt_proj, '(b c) d -> b c d', b=b)

        loss_finegrained = trainer.entity_grounder_loss(image_proj, txt_proj, trainer.text_model.tau, mask)
        trainer.backward(loss_finegrained)
        train_loss += loss_finegrained.item()

        pbar.update(1)
        index += 1
        pbar.set_description(f"epoch {trainer.epoch + 1} Entity Grounder Loss: {train_loss / index}")
        trainer.log({"entity_grounder_loss": loss_finegrained.item()})
        trainer.train_steps += 1
    trainer.epoch += 1


def trainSegmentMerger(data, model, criterion, optimizer, epoch):  # todo: refactor this function to use trainer class
    train_loss = 0.0
    pbar = tqdm.tqdm(data, total=len(data))
    index = 0
    for image_segs, text_segs, images, strings in data:
        if image_segs.size(0) != BATCH_SIZE:
            continue
        image_segs = image_segs.detach().to(device)
        text_segs = text_segs.detach().to(device)
        image_proj = model(image_segs, None)

        b = image_proj.size(0)
        # Assuming that the padding value is 0, generate a mask for non-padded elements

        text_segs = rearrange(text_segs, 'b c t d -> (b c) t d')
        txt_proj = model(None, text_segs)
        txt_proj = rearrange(txt_proj, '(b c) d -> b c d', b=b)
        txt_proj = txt_proj.detach().to(device)

        loss, sims = criterion(image_proj, txt_proj)
        optimizer.zero_grad()

        loss.backward()
        optimizer.step()
        train_loss += loss.item()
        pbar.update(1)
        index += 1
        pbar.set_description(f"epoch {epoch + 1} Segment Merger --- Loss: {train_loss / index}")
        wandb.log({"segment_merger_loss": loss.item()})

    return model


def main():
    """
    The main function initializes the models, datasets, optimizers, and starts the training loop.
    
    Args:
    None
    
    Returns:
    None
    """
    print("done processing dataset")

    trainer = Trainer(settings)
    trainer.setup()
    for epoch in range(100):
        trainEntityGrounder(trainer)
        verifyEntityGrounder(trainer, 1)
        trainer.save()


main()