Spaces:

alvanlii
/

FROMAGe

Runtime error

FROMAGe / main.py

alvanli

Add cheese model

1f43fd8 over 1 year ago

No virus

29.7 kB

	"""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:])