train_fashion200k.py

import argparse
import os
import ruamel.yaml as yaml
import numpy as np
import random
import time
import datetime
import json
from pathlib import Path
import cv2
from tqdm import tqdm

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
import torch.distributed as dist
from torch.utils.data import DataLoader
from torch import optim


from data import create_dataset, create_sampler, create_loader
from models.combiner import CombinerModel
from models.blip_itm import blip_itm
from utils import update_train_running_results, set_train_bar_description


def main(args, config):
    device = torch.device(args.device)
    seed = 42

    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    cudnn.benchmark = True

    #### Dataset ####
    print("Creating fashion200k dataset")
    train_dataset, test_dataset = create_dataset('fashion200k', config)

    samplers = [None, None]
    train_loader, test_loader = create_loader([train_dataset, test_dataset], samplers,
                                              batch_size=[
                                                  config['batch_size_train']]+[config['batch_size_test']],
                                              num_workers=[4, 4],
                                              is_trains=[True, False],
                                              collate_fns=[None, None])

    # init combiner
    combiner = CombinerModel(
        config['v_dim'], config['l_dim'], config['dim'], config['num_heads'])
    combiner = combiner.to(device)
    # Define the optimizer, the loss and the grad scaler
    optimizer = optim.Adam(combiner.parameters(), lr=config['combiner_lr'])
    crossentropy_criterion = nn.CrossEntropyLoss()
    scaler = torch.cuda.amp.GradScaler()

    # init BLIP pretrained model to use their encoders
    print('loading pretrained BLIP')
    blip = blip_itm(pretrained=config['model_url'], image_size=config['image_size'], vit = 'base')
    blip = blip.to(device)
    print('BLIP loaded succesfuly')

    print('========== Start training loop ========== ')

    for epoch in range(config['num_epochs']):
        if torch.cuda.is_available():
            combiner.train()
            train_running_results = {
                'images_in_epoch': 0, 'accumulated_train_loss': 0}
            train_bar = tqdm(train_loader, ncols=150)

            for idx, (out) in enumerate(train_bar):  # Load a batch of data

                reference_images = out['source_img_data']
                target_images = out['target_img_data']
                captions = out['mod']

                images_in_batch = reference_images.size(0)

                optimizer.zero_grad()
                step = len(train_bar) * epoch + idx

                reference_images = reference_images.to(
                    device, non_blocking=True)
                target_images = target_images.to(device, non_blocking=True)

                input_captions: list = np.array(
                    captions).T.flatten().tolist()

                text_inputs = blip.tokenizer(input_captions, padding='max_length', truncation=True, max_length=config['max_length'], return_tensors="pt").to(
                    device)  # FIXME  double check if correct

                # Extract the features with BLIP here
                with torch.no_grad():
                    # extract image features
                    reference_image_features = blip.visual_encoder(reference_images)
                    target_image_features = blip.visual_encoder(target_images)
                    # extract text features
                    text_features = blip.text_encoder(text_inputs.input_ids, attention_mask=text_inputs.attention_mask,
                                                      return_dict=True, mode='text')

                    # text_features = {last_hidden_state, pooler_output}
                    # https://github.com/huggingface/transformers/issues/7540#issuecomment-704155218

                with torch.cuda.amp.autocast():
                    # feed extracted features into combiner
                    combiner_out_v, combiner_out_l = combiner(reference_image_features,
                                                              text_features['last_hidden_state'], text_inputs.attention_mask[0, :])
                    # cls tokens are the first token in the sequence (for both vision & text), use a projection layer to ensure their dimensions are the same
                    image_feat = F.normalize(combiner.vision_proj(combiner_out_v[:, 0, :]), dim=-1)
                    text_feat = F.normalize(combiner.text_proj(combiner_out_l[:, 0, :]), dim=-1)
                    target_image_feat = F.normalize(combiner.vision_proj(target_image_features[:, 0, :]), dim=-1)
                    # calculate cosine similarity, which is used for contrastive loss
                    sim = image_feat @ target_image_feat.t()
                    ground_truth = torch.arange(images_in_batch, dtype=torch.long, device=device)

                    loss = crossentropy_criterion(sim, ground_truth)
                    break

                # Backpropagate and update the weights
                scaler.scale(loss).backward()
                scaler.step(optimizer)
                scaler.update()

                update_train_running_results(
                    train_running_results, loss, images_in_batch)
                set_train_bar_description(
                    train_bar, epoch, config['num_epochs'], train_running_results)

            train_epoch_loss = float(
                train_running_results['accumulated_train_loss'] / train_running_results['images_in_epoch'])


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', default='./configs/fashion200k.yaml')
    parser.add_argument('--device', default='cuda')

    args = parser.parse_args()
    config = yaml.load(open(args.config, 'r'), Loader=yaml.Loader)

    main(args, config)