main.py

"""Training example.

Modified from https://github.com/pytorch/examples/blob/main/imagenet/main.py.
"""
import argparse
import json
import os
import sys
import time
import warnings

import numpy as np
from PIL import Image
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import torch.optim
from torch.optim.lr_scheduler import StepLR
from warmup_scheduler import GradualWarmupScheduler
import torch.multiprocessing as mp
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch.utils.tensorboard import SummaryWriter
import torchvision

from fromage import data
from fromage import losses as losses_utils
from fromage import models
from fromage import utils
from fromage import evaluate
from transformers import AutoTokenizer

# Disable HuggingFace tokenizer parallelism.
os.environ["TOKENIZERS_PARALLELISM"] = "false"

# Available LLM models.
llm_models = ['facebook/opt-125m', 'facebook/opt-350m', 'facebook/opt-1.3b',
              'facebook/opt-2.7b', 'facebook/opt-6.7b', 'facebook/opt-13b', 'facebook/opt-30b',
              'facebook/opt-66b']
datasets = ['cc3m']
best_score = 0  # Variable to keep track of best model so far.


def parse_args(args):
  parser = argparse.ArgumentParser(description='FROMAGe training')
  parser.add_argument('--opt-version', default='facebook/opt-6.7b',
            choices=llm_models,
            help='OPT versions: ' +
              ' | '.join(llm_models) +
              ' (default: "facebook/opt-6.7b")')
  parser.add_argument('--visual-model', default='openai/clip-vit-large-patch14', type=str,
                      help="Visual encoder to use.")
  parser.add_argument('-d', '--dataset', metavar='DATASET',  help='Delimited list of datasets:' +
                      ' | '.join(datasets), default='cc3m',
                      type=lambda s: [x for x in s.split(',')])

  parser.add_argument('--val-dataset', metavar='DATASET', default='cc3m',
            type=lambda s: [x for x in s.split(',')],
            help='Validation dataset: ' +
              ' | '.join(datasets) +
              ' (default: cc3m)')
  parser.add_argument('--dataset_dir', default='datasets', type=str,
            help='Dataset directory containing .tsv files.')
  parser.add_argument('--image-dir', default='./data/', type=str,
            help='Dataset directory containing image folders.')
  parser.add_argument('--log-base-dir', default='./runs/', type=str,
            help='Base directory to write logs and ckpts to.')
  parser.add_argument('--exp_name', default='frozen', type=str,
            help='Name of experiment, used for saving checkpoints.')

  parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
            help='number of data loading workers (default: 4)')
  parser.add_argument('--epochs', default=10, type=int, metavar='N',
            help='number of total epochs to run')
  parser.add_argument('--steps-per-epoch', default=2000, type=int, metavar='N',
            help='number of training steps per epoch')
  parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
            help='manual epoch number (useful on restarts)')
  parser.add_argument('--val-steps-per-epoch', default=-1, type=int, metavar='N',
            help='number of validation steps per epoch.')
  parser.add_argument('-b', '--batch-size', default=180, type=int,
            metavar='N',
            help='mini-batch size (default: 180), this is the total '
               'batch size of all GPUs on the current node when '
               'using Data Parallel or Distributed Data Parallel')
  parser.add_argument('--val-batch-size', default=None, type=int)
  parser.add_argument('--lr', '--learning-rate', default=0.0003, type=float,
            metavar='LR', help='initial learning rate', dest='lr')
  parser.add_argument('--lr-warmup-steps', default=100, type=int,
            metavar='N', help='Number of steps to warm up lr.')
  parser.add_argument('--lr-schedule-step-size', default=10, type=int,
            metavar='N', help='Number of steps before decaying lr.')
  parser.add_argument('--lr-schedule-gamma', default=0.1, type=float,
            metavar='N', help='Decay parameter for learning rate scheduler.')
  parser.add_argument('--grad-accumulation-steps', default=1, type=int, metavar='N',
                      help='number of gradient accumulation steps')
  parser.add_argument('--grad-clip', default=1.0, type=float, help='gradient clipping amount')

  parser.add_argument('--precision', default='fp32', type=str, choices=['fp32', 'fp16', 'bf16'], help="Precision to train in.")
  parser.add_argument('--cap-loss-scale', type=float, default=1.0, help="Scale on captioning loss.")
  parser.add_argument('--ret-loss-scale', type=float, default=1.0, help="Scale on retrieval loss.")

  parser.add_argument('--concat-captions-prob', type=float, default=0.5, help="Probability of concatenating two examples sequentially for captioning.")
  parser.add_argument('--concat-for-ret', action='store_true', default=False, help="Whether to concatenate examples for retrieval mode.")
  parser.add_argument('--input-prompt', default=None, type=str, help="Input prompt for the language model, if any.")

  parser.add_argument('--image-size', default=224, type=int, metavar='N', help='Size of images.')
  parser.add_argument('--use_image_embed_norm', action='store_true', default=False, help="Whether to use norm on the image embeddings to make them equal to language.")
  parser.add_argument('--image_embed_dropout_prob', type=float, default=0.0, help="Dropout probability on the image embeddings.")
  parser.add_argument('--use_text_embed_layernorm', action='store_true', default=False, help="Whether to use layer norm on the text embeddings for retrieval.")
  parser.add_argument('--text_embed_dropout_prob', type=float, default=0.0, help="Dropout probability on the text embeddings.")
  parser.add_argument('--shared-emb-dim', default=256, type=int, metavar='N', help='Embedding dimension for retrieval.')
  parser.add_argument('--text-emb-layers', help='Layer to use for text embeddings. OPT-2.7b has 33 layers.', default='-1',
                      type=lambda s: [int(x) for x in s.split(',')])

  parser.add_argument('--max-len', default=24, type=int,
            metavar='N', help='Maximum length to truncate captions / generations to.')
  parser.add_argument('--n-visual-tokens', default=1, type=int,
            metavar='N', help='Number of visual tokens to use for the Frozen model.')

  parser.add_argument('--beta1', default=0.9, type=float, metavar='M', help='beta1 for Adam')
  parser.add_argument('--beta2', default=0.95, type=float, metavar='M', help='beta2 for Adam')
  parser.add_argument('--wd', '--weight-decay', default=0.0, type=float,
            metavar='W', help='weight decay (default: 0.0)', dest='weight_decay')
  parser.add_argument('-p', '--print-freq', default=10, type=int,
            metavar='N', help='print frequency (default: 10)')
  parser.add_argument('--resume', default='', type=str, metavar='PATH',
            help='path to latest checkpoint (default: none)')
  parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',
            help='evaluate model on validation set')
  parser.add_argument('--world-size', default=-1, type=int,
            help='number of nodes for distributed training')
  parser.add_argument('--rank', default=-1, type=int,
            help='node rank for distributed training')
  parser.add_argument('--dist-url', default='tcp://127.0.0.1:1337', type=str,
            help='url used to set up distributed training')
  parser.add_argument('--dist-backend', default='nccl', type=str,
            help='distributed backend')
  parser.add_argument('--seed', default=None, type=int,
            help='seed for initializing training. ')
  parser.add_argument('--gpu', default=None, type=int,
            help='GPU id to use.')
  parser.add_argument('--multiprocessing-distributed', action='store_true',
            help='Use multi-processing distributed training to launch '
               'N processes per node, which has N GPUs. This is the '
               'fastest way to use PyTorch for either single node or '
               'multi node data parallel training')
  return parser.parse_args(args)


def main(args):
  args = parse_args(args)
  i = 1
  args.log_dir = os.path.join(args.log_base_dir, args.exp_name)
  while os.path.exists(args.log_dir):
    args.log_dir = os.path.join(args.log_base_dir, f'{args.exp_name}_{i}')
    i += 1
  os.makedirs(args.log_dir)

  with open(os.path.join(args.log_dir, f'args.json'), 'w') as wf:
    json.dump(vars(args), wf, indent=4)

  with open(os.path.join(args.log_dir, f'git_info.txt'), 'w') as wf:
    utils.dump_git_status(out_file=wf)

  print(f'Logging to {args.log_dir}.')

  if args.seed is not None:
    torch.manual_seed(args.seed)
    cudnn.deterministic = True
    warnings.warn('You have chosen to seed training. '
            'This will turn on the CUDNN deterministic setting, '
            'which can slow down your training considerably! '
            'You may see unexpected behavior when restarting '
            'from checkpoints.')

  if args.gpu is not None:
    warnings.warn('You have chosen a specific GPU. This will completely '
            'disable data parallelism.')

  if args.dist_url == "env://" and args.world_size == -1:
    args.world_size = int(os.environ["WORLD_SIZE"])

  args.distributed = args.world_size > 1 or args.multiprocessing_distributed

  ngpus_per_node = torch.cuda.device_count()
  if args.multiprocessing_distributed:
    # Since we have ngpus_per_node processes per node, the total world_size
    # needs to be adjusted accordingly
    args.world_size = ngpus_per_node * args.world_size
    # Use torch.multiprocessing.spawn to launch distributed processes: the
    # main_worker process function
    mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args))
  else:
    # Simply call main_worker function
    main_worker(args.gpu, ngpus_per_node, args)


def main_worker(gpu, ngpus_per_node, args):
  """Setup code."""
  global best_score
  args.gpu = gpu

  if args.gpu is not None:
    print("Use GPU: {} for training".format(args.gpu))

  if args.distributed:
    if args.dist_url == "env://" and args.rank == -1:
      args.rank = int(os.environ["RANK"])
    if args.multiprocessing_distributed:
      # For multiprocessing distributed training, rank needs to be the
      # global rank among all the processes
      args.rank = args.rank * ngpus_per_node + gpu
    dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
                world_size=args.world_size, rank=args.rank)

  # Create model
  model_args = models.FrozenArgs()
  model_args.opt_version = args.opt_version
  model_args.freeze_lm = True
  model_args.visual_encoder = args.visual_model
  model_args.freeze_vm = True
  model_args.n_visual_tokens = args.n_visual_tokens
  model_args.use_image_embed_norm = args.use_image_embed_norm
  model_args.image_embed_dropout_prob = args.image_embed_dropout_prob
  model_args.use_text_embed_layernorm = args.use_text_embed_layernorm
  model_args.text_embed_dropout_prob = args.text_embed_dropout_prob
  model_args.shared_emb_dim = args.shared_emb_dim
  model_args.text_emb_layers = args.text_emb_layers

  tokenizer = AutoTokenizer.from_pretrained(args.opt_version, use_fast=False)
  # Add an image token for loss masking (and visualization) purposes.
  tokenizer.add_special_tokens({"cls_token": "<|image|>"})  # add special image token to tokenizer
  print('Adding [RET] token to vocabulary.')
  print('Before adding new token, tokenizer("[RET]") =', tokenizer('[RET]', add_special_tokens=False))
  num_added_tokens = tokenizer.add_tokens('[RET]')
  print(f'After adding {num_added_tokens} new tokens, tokenizer("[RET]") =', tokenizer('[RET]', add_special_tokens=False))
  ret_token_idx = tokenizer('[RET]', add_special_tokens=False).input_ids
  assert len(ret_token_idx) == 1, ret_token_idx
  model_args.retrieval_token_idx = ret_token_idx[0]
  args.retrieval_token_idx = ret_token_idx[0]

  # Save model args to disk.
  with open(os.path.join(args.log_dir, 'model_args.json'), 'w') as f:
    json.dump(vars(model_args), f, indent=4)

  model = models.Fromage(tokenizer, model_args)
  if args.precision == 'fp16':
    model = model.float()
  elif args.precision == 'bf16':
    model = model.bfloat16()

  # Print parameters and count of model.
  param_counts_text = utils.get_params_count_str(model)
  with open(os.path.join(args.log_dir, 'param_count.txt'), 'w') as f:
    f.write(param_counts_text)

  # Log trainable parameters to Tensorboard.
  _, total_trainable_params, total_nontrainable_params = utils.get_params_count(model)
  writer = SummaryWriter(args.log_dir)
  writer.add_scalar('params/total', total_trainable_params + total_nontrainable_params, 0)
  writer.add_scalar('params/total_trainable', total_trainable_params, 0)
  writer.add_scalar('params/total_non_trainable', total_nontrainable_params, 0)
  writer.close()

  if not torch.cuda.is_available():
    print('WARNING: using CPU, this will be slow!')
    model = torch.nn.DataParallel(model)
  elif args.distributed:
    # For multiprocessing distributed, DistributedDataParallel constructor
    # should always set the single device scope, otherwise,
    # DistributedDataParallel will use all available devices.
    if args.gpu is not None:
      torch.cuda.set_device(args.gpu)
      model.cuda(args.gpu)
      # When using a single GPU per process and per
      # DistributedDataParallel, we need to divide the batch size
      # ourselves based on the total number of GPUs of the current node.
      args.batch_size = int(args.batch_size / ngpus_per_node)
      args.val_batch_size = int((args.val_batch_size or args.batch_size) / ngpus_per_node)
      args.workers = int((args.workers + ngpus_per_node - 1) / ngpus_per_node)
      model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=False)
    else:
      model.cuda()
      # DistributedDataParallel will divide and allocate batch_size to all
      # available GPUs if device_ids are not set
      model = torch.nn.parallel.DistributedDataParallel(model, find_unused_parameters=False)
  elif args.gpu is not None:
    torch.cuda.set_device(args.gpu)
    model = model.cuda(args.gpu)
  else:
    model = torch.nn.DataParallel(model).cuda()

  # define loss function (criterion), optimizer, and learning rate scheduler
  criterion = nn.CrossEntropyLoss().cuda(args.gpu)
  optimizer_cls = torch.optim.AdamW
  print('Using torch.optim.AdamW as the optimizer.')
  optimizer = optimizer_cls(model.parameters(), args.lr,
                betas=(args.beta1, args.beta2),
                weight_decay=args.weight_decay,
                eps=1e-8)

  """Sets the learning rate to the initial LR decayed by 10 every 5 epochs"""
  scheduler_steplr = StepLR(optimizer, step_size=args.lr_schedule_step_size * args.steps_per_epoch, gamma=args.lr_schedule_gamma)
  scheduler = GradualWarmupScheduler(optimizer, multiplier=1.0, total_epoch=args.lr_warmup_steps, after_scheduler=scheduler_steplr)
  
  # optionally resume from a checkpoint
  if args.resume:
    if os.path.isfile(args.resume):
      print("=> loading checkpoint '{}'".format(args.resume))
      if args.gpu is None:
        checkpoint = torch.load(args.resume)
      else:
        # Map model to be loaded to specified single gpu.
        loc = 'cuda:{}'.format(args.gpu)
        checkpoint = torch.load(args.resume, map_location=loc)
      args.start_epoch = checkpoint['epoch']
      best_score = checkpoint['best_score']
      if args.gpu is not None:
        # best_score may be from a checkpoint from a different GPU
        best_score = best_score.to(args.gpu)
      model.load_state_dict(checkpoint['state_dict'])
      optimizer.load_state_dict(checkpoint['optimizer'])
      scheduler.load_state_dict(checkpoint['scheduler'])
      print("=> loaded checkpoint '{}' (epoch {})"
          .format(args.resume, checkpoint['epoch']))
    else:
      print("=> no checkpoint found at '{}'".format(args.resume))

  cudnn.benchmark = True

  # Data loading code
  train_dataset = data.get_dataset(args, 'train', tokenizer)
  val_dataset = data.get_dataset(args, 'val', tokenizer)
  print(f'Training with {len(train_dataset)} examples and validating with {len(val_dataset)} examples.')

  if args.distributed:
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset, drop_last=True)
    val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset, shuffle=False, drop_last=True)
  else:
    train_sampler = None
    val_sampler = None

  train_loader = torch.utils.data.DataLoader(
    train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),
    num_workers=args.workers, pin_memory=True, sampler=train_sampler)
  val_loader = torch.utils.data.DataLoader(
    val_dataset, batch_size=(args.val_batch_size or args.batch_size), shuffle=False,
    num_workers=args.workers, pin_memory=True, sampler=val_sampler)

  if args.evaluate:
    evaluate.validate(val_loader, model, tokenizer, criterion, epoch, args)
    return

  for epoch in range(args.start_epoch, args.epochs):
    if epoch == 0:
      evaluate.validate(val_loader, model, tokenizer, criterion, epoch-1, args)
    if args.distributed:
      train_sampler.set_epoch(epoch)

    # train for one epoch
    train(train_loader, model, tokenizer, criterion, optimizer, epoch, scheduler, args)

    # evaluate on validation set
    eval_score = evaluate.validate(val_loader, model, tokenizer, criterion, epoch, args)

    # remember best score and save checkpoint
    is_best = eval_score > best_score
    best_score = max(eval_score, best_score)

    if not args.multiprocessing_distributed or (args.multiprocessing_distributed
        and args.rank % ngpus_per_node == 0):
      utils.save_checkpoint({
        'epoch': epoch + 1,
        'state_dict': model.state_dict(),
        'best_score': best_score,
        'optimizer' : optimizer.state_dict(),
        'scheduler' : scheduler.state_dict()
      }, is_best, os.path.join(args.log_dir, 'ckpt'))


def train(train_loader, model, tokenizer, criterion, optimizer, epoch, scheduler, args):
  """Main training loop."""
  ngpus_per_node = torch.cuda.device_count()
  batch_time = utils.AverageMeter('Time', ':6.3f')
  cap_time = utils.AverageMeter('CaptioningTime', ':6.3f')
  ret_time = utils.AverageMeter('RetrievalTime', ':6.3f')
  data_time = utils.AverageMeter('Data', ':6.3f')
  losses = utils.AverageMeter('Loss', ':.4e')
  ce_losses = utils.AverageMeter('CeLoss', ':.4e')
  top1 = utils.AverageMeter('Acc@1', ':6.2f')
  top5 = utils.AverageMeter('Acc@5', ':6.2f')
  cont_losses = utils.AverageMeter('ContLoss', ':.4e')
  top1_caption = utils.AverageMeter('AccCaption@1', ':6.2f')
  top5_caption = utils.AverageMeter('AccCaption@5', ':6.2f')
  top1_image = utils.AverageMeter('AccImage@1', ':6.2f')
  top5_image = utils.AverageMeter('AccImage@5', ':6.2f')

  writer = SummaryWriter(args.log_dir)

  progress = utils.ProgressMeter(
    args.steps_per_epoch,
    [batch_time, losses, ce_losses, cont_losses, top1, top5],
    prefix="Epoch: [{}]".format(epoch))

  # switch to train mode
  model.train()

  end = time.time()

  for i, (image_paths, images, caption_images, tgt_tokens, token_len) in enumerate(train_loader):
    actual_step = epoch * args.steps_per_epoch + i + 1
    # measure data loading time
    data_time.update(time.time() - end)

    if torch.cuda.is_available():
      images = images.cuda(args.gpu, non_blocking=True)
      tgt_tokens = tgt_tokens.cuda(args.gpu, non_blocking=True)
      token_len = token_len.cuda(args.gpu, non_blocking=True)

    if args.precision == 'fp16':
      images = images.half()
    elif args.precision == 'bf16':
      images = images.bfloat16()

    model_modes = ['captioning', 'retrieval']
    loss = 0

    for model_mode in model_modes:
      mode_start = time.time()
      # compute output
      concat_captions = np.random.uniform(0, 1) < args.concat_captions_prob
      if not args.concat_for_ret:
        concat_captions = concat_captions and model_mode == 'captioning'

      (model_output, full_labels, last_embedding, _, visual_embs) = model(
        images, tgt_tokens, token_len, mode=model_mode, concat_captions=concat_captions, inference=False)
      output = model_output.logits

      # Measure captioning accuracy for multi-task models and next-token prediction for retrieval models.
      if model_mode == 'captioning':
        acc1, acc5 = utils.accuracy(output[:, :-1, :], full_labels[:, 1:], -100, topk=(1, 5))
        top1.update(acc1[0], images.size(0))
        top5.update(acc5[0], images.size(0))

      ce_loss = model_output.loss
      if model_mode == 'captioning':
        ce_loss = ce_loss * args.cap_loss_scale
      elif model_mode == 'retrieval':
        ce_loss = ce_loss * args.ret_loss_scale
      else:
        raise NotImplementedError

      loss += ce_loss
      ce_losses.update(ce_loss.item(), images.size(0))

      if model_mode == 'retrieval':
        # Cross replica concat for embeddings.
        if args.distributed:
          all_visual_embs = [torch.zeros_like(visual_embs) for _ in range(dist.get_world_size())]
          all_last_embedding = [torch.zeros_like(last_embedding) for _ in range(dist.get_world_size())]
          dist.all_gather(all_visual_embs, visual_embs)
          dist.all_gather(all_last_embedding, last_embedding)
          # Overwrite with embeddings produced on this replace, which have the gradient.
          all_visual_embs[dist.get_rank()] = visual_embs
          all_last_embedding[dist.get_rank()] = last_embedding
          visual_embs = torch.cat(all_visual_embs)
          last_embedding = torch.cat(all_last_embedding)

          start_idx = args.rank * images.shape[0]
          end_idx = start_idx + images.shape[0]

        logits_per_image = visual_embs @ last_embedding.t()
        logits_per_text = logits_per_image.t()
        if i == 0:
          print(f'Running contrastive loss over logits_per_text.shape = {logits_per_text.shape} and logits_per_image.shape = {logits_per_image.shape}')

        # Compute contrastive losses for retrieval.
        caption_loss = losses_utils.contrastive_loss(logits_per_text)
        image_loss = losses_utils.contrastive_loss(logits_per_image)
        caption_acc1, caption_acc5 = losses_utils.contrastive_acc(logits_per_text, topk=(1, 5))
        image_acc1, image_acc5 = losses_utils.contrastive_acc(logits_per_image, topk=(1, 5))
        loss += args.ret_loss_scale * (caption_loss + image_loss) / 2.0
        cont_losses.update(loss.item(), images.size(0))

        # measure accuracy and record loss
        top1_caption.update(caption_acc1[0], images.size(0))
        top5_caption.update(caption_acc5[0], images.size(0))
        top1_image.update(image_acc1[0], images.size(0))
        top5_image.update(image_acc5[0], images.size(0))

      if model_mode == 'retrieval':
        ret_time.update(time.time() - mode_start)
      elif model_mode == 'captioning':
        cap_time.update(time.time() - mode_start)

    loss = loss / args.grad_accumulation_steps
    losses.update(loss.item(), images.size(0))
    loss.backward()

    # Update weights
    if ((i + 1) % args.grad_accumulation_steps == 0) or (i == args.steps_per_epoch - 1):
      # Zero out gradients of the embedding matrix outside of [RET].
      for param in model.module.model.input_embeddings.parameters():
        assert param.grad.shape[0] == len(tokenizer)
        # Keep other embeddings frozen.
        mask = torch.arange(param.grad.shape[0]) != args.retrieval_token_idx
        param.grad[mask, :] = 0

      # compute gradient and do SGD step
      if args.grad_clip > 0:
        nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
      optimizer.step()
      optimizer.zero_grad()

    with torch.no_grad():
      # Normalize trainable embeddings.
      frozen_norm = torch.norm(model.module.model.input_embeddings.weight[:-1, :], dim=1).mean(0)
      trainable_weight = model.module.model.input_embeddings.weight[-1, :]
      model.module.model.input_embeddings.weight[-1, :].div_(torch.norm(trainable_weight) / frozen_norm)

    # measure elapsed time
    batch_time.update(time.time() - end)
    end = time.time()

    if actual_step == 1 or (i + 1) % args.print_freq == 0:
      ex_per_sec = args.batch_size / batch_time.avg
      if args.distributed:
        batch_time.all_reduce()
        data_time.all_reduce()
        ex_per_sec = (args.batch_size / batch_time.avg) * ngpus_per_node

        losses.all_reduce()
        ce_losses.all_reduce()
        top1.all_reduce()
        top5.all_reduce()
        ret_time.all_reduce()
        cont_losses.all_reduce()
        top1_caption.all_reduce()
        top5_caption.all_reduce()
        top1_image.all_reduce()
        top5_image.all_reduce()
        cap_time.all_reduce()

      progress.display(i + 1)

      writer.add_scalar('train/loss', losses.avg, actual_step)
      writer.add_scalar('train/ce_loss', ce_losses.avg, actual_step)
      writer.add_scalar('train/seq_top1_acc', top1.avg, actual_step)
      writer.add_scalar('train/seq_top5_acc', top5.avg, actual_step)
      writer.add_scalar('train/contrastive_loss', cont_losses.avg, actual_step)
      writer.add_scalar('train/t2i_top1_acc', top1_caption.avg, actual_step)
      writer.add_scalar('train/t2i_top5_acc', top5_caption.avg, actual_step)
      writer.add_scalar('train/i2t_top1_acc', top1_image.avg, actual_step)
      writer.add_scalar('train/i2t_top5_acc', top5_image.avg, actual_step)
      writer.add_scalar('metrics/total_secs_per_batch', batch_time.avg, actual_step)
      writer.add_scalar('metrics/total_secs_captioning', cap_time.avg, actual_step)
      writer.add_scalar('metrics/total_secs_retrieval', ret_time.avg, actual_step)
      writer.add_scalar('metrics/data_secs_per_batch', data_time.avg, actual_step)
      writer.add_scalar('metrics/examples_per_sec', ex_per_sec, actual_step)

      if not args.multiprocessing_distributed or (args.multiprocessing_distributed
        and args.rank % ngpus_per_node == 0):
        image_bs = images.shape[0]
        normalized_images = images - images.min()
        normalized_images /= normalized_images.max()  # (N, 3, H, W)
        max_images_to_show = 16

        # Append caption text.
        pred_tokens = output[:, args.n_visual_tokens-1:-1, :].argmax(dim=-1)
        generated_captions = tokenizer.batch_decode(pred_tokens, skip_special_tokens=False)

        # Log image (and generated caption) outputs to Tensorboard.
        if model_mode == 'captioning':
          # Create generated caption text.
          generated_cap_images = torch.stack([
            utils.create_image_of_text(
              generated_captions[i].encode('ascii', 'ignore'),
              width=normalized_images.shape[3],
              color=(255, 255, 0))
            for i in range(len(generated_captions))], axis=0)

          # Duplicate captions if we concatenated them.
          if (args.concat_captions_prob > 0 and model_mode == 'captioning' and generated_cap_images.shape[0] != caption_images.shape[0]):
            generated_cap_images = torch.cat([generated_cap_images, generated_cap_images], axis=0)

          display_images = torch.cat([normalized_images.float().cpu(), caption_images, generated_cap_images], axis=2)[:max_images_to_show]
          grid = torchvision.utils.make_grid(display_images, nrow=int(max_images_to_show ** 0.5), padding=4)
          writer.add_image('train/images_gen_cap', grid, actual_step)

        # Retrieved images (from text).
        retrieved_image_idx = logits_per_text[:image_bs, :image_bs].argmax(-1)
        t2i_images = torch.stack(
          [normalized_images[retrieved_image_idx[i], ...] for i in range(len(retrieved_image_idx))],
          axis=0)
        t2i_images = torch.cat([t2i_images.float().cpu(), caption_images], axis=2)[:max_images_to_show]
        t2i_grid = torchvision.utils.make_grid(t2i_images, nrow=int(max_images_to_show ** 0.5), padding=4)
        writer.add_image('train/t2i_ret', t2i_grid, actual_step)

        # Retrieved text (from image).
        retrieved_text_idx = logits_per_image[:image_bs, :image_bs].argmax(-1)
        retrieved_text = torch.stack(
          [caption_images[retrieved_text_idx[i], ...] for i in range(len(retrieved_text_idx))],
          axis=0)
        i2t_images = torch.cat([normalized_images.float().cpu(), retrieved_text], axis=2)[:max_images_to_show]
        i2t_grid = torchvision.utils.make_grid(i2t_images, nrow=int(max_images_to_show ** 0.5), padding=4)
        writer.add_image('train/i2t_ret', i2t_grid, actual_step)

      batch_time.reset()
      cap_time.reset()
      ret_time.reset()
      data_time.reset()
      losses.reset()
      ce_losses.reset()
      top1.reset()
      top5.reset()
      cont_losses.reset()
      top1_caption.reset()
      top5_caption.reset()
      top1_image.reset()
      top5_image.reset()

    if i == args.steps_per_epoch - 1:
      break

    scheduler.step()
    curr_lr = scheduler.get_last_lr()
    if (actual_step == 1) or (i + 1) % args.print_freq == 0:
      # Write current learning rate to Tensorboard.
      writer = SummaryWriter(args.log_dir)
      writer.add_scalar('train/lr', curr_lr[0], actual_step)
      writer.close()

  writer.close()


if __name__ == '__main__':
  main(sys.argv[1:])