Add files using upload-large-folder tool

6b56858 verified 5 days ago

11.3 kB

import json
import os

import fsspec
import hydra
import lightning as L
from lightning.fabric import Fabric
import omegaconf
import rich.syntax
import rich.tree
import torch
from torch.utils.data.distributed import DistributedSampler
from torchmetrics.image.fid import FrechetInceptionDistance
from torchmetrics.image.inception import InceptionScore
from tqdm import tqdm, trange

import algo
import dataloader
import utils

omegaconf.OmegaConf.register_new_resolver(
  'cwd', os.getcwd)
omegaconf.OmegaConf.register_new_resolver(
  'device_count', torch.cuda.device_count)
omegaconf.OmegaConf.register_new_resolver(
  'eval', eval)
omegaconf.OmegaConf.register_new_resolver(
  'div_up', lambda x, y: (x + y - 1) // y)


def _load_from_checkpoint(diffusion_model, config, tokenizer):
  if 'hf' in config.algo.backbone:
    return diffusion_model(
      config, tokenizer=tokenizer).to('cuda')
  
  return diffusion_model.load_from_checkpoint(
    config.eval.checkpoint_path,
    tokenizer=tokenizer,
    config=config)


@L.pytorch.utilities.rank_zero_only
def _print_config(
  config: omegaconf.DictConfig,
  resolve: bool = True,
  save_cfg: bool = True) -> None:
  """Prints content of DictConfig using Rich library and its tree structure.
  
  Args:
    config (DictConfig): Configuration composed by Hydra.
    resolve (bool): Whether to resolve reference fields of DictConfig.
    save_cfg (bool): Whether to save the configuration tree to a file.
  """

  style = 'dim'
  tree = rich.tree.Tree('CONFIG', style=style, guide_style=style)

  fields = config.keys()
  for field in fields:
    branch = tree.add(field, style=style, guide_style=style)

    config_section = config.get(field)
    branch_content = str(config_section)
    if isinstance(config_section, omegaconf.DictConfig):
      branch_content = omegaconf.OmegaConf.to_yaml(
        config_section, resolve=resolve)

    branch.add(rich.syntax.Syntax(branch_content, 'yaml'))
  rich.print(tree)
  if save_cfg:
    with fsspec.open(
      '{}/config_tree.txt'.format(
        config.checkpointing.save_dir), 'w') as fp:
      rich.print(tree, file=fp)


@L.pytorch.utilities.rank_zero_only
def _print_batch(config, train_ds, valid_ds, tokenizer, k=64):
  for dl_type, dl in [
    ('train', train_ds), ('valid', valid_ds)]:
    print(f'Printing {dl_type} dataloader batch.')
    batch = next(iter(dl))
    print('Batch input_ids.shape', batch['input_ids'].shape)
    if config.data.modality == 'text':
      first = batch['input_ids'][0, :k]
      last = batch['input_ids'][0, -k:]
      print(f'First {k} tokens:', tokenizer.decode(first))
      print('ids:', first)
      print(f'Last {k} tokens:', tokenizer.decode(last))
      print('ids:', last)


def _generate_samples(diffusion_model, config, logger,
                      tokenizer):
  logger.info('Starting Sample Eval.')
  model = _load_from_checkpoint(
    diffusion_model=diffusion_model,
    config=config,
    tokenizer=tokenizer)
  model.metrics.gen_ppl.reset()
  model.metrics.sample_entropy.reset()
  if config.eval.disable_ema:
    logger.info('Disabling EMA.')
    model.ema = None
  stride_length = config.sampling.stride_length
  num_strides = config.sampling.num_strides
  all_samples = []
  for _ in trange(config.sampling.num_sample_batches):
    if config.sampling.semi_ar:
      _, intermediate_samples, _ = model.restore_model_and_semi_ar_sample(
        stride_length=stride_length,
        num_strides=num_strides,
        dt=1 / config.sampling.steps)
      text_samples = intermediate_samples[-1]
      # Note: Samples generated using semi-ar method
      # need to to be processed before computing generative perplexity
      # since these samples contain numerous <|endoftext|> tokens
      # and diffusion.compute_generative_perplexity() discards
      # any text after the first EOS token.
    else:
      samples = model.restore_model_and_sample(
        num_steps=config.sampling.steps)
      model.metrics.record_entropy(samples)
      text_samples = model.tokenizer.batch_decode(samples)
      model.metrics.record_generative_perplexity(
        text_samples, config.model.length, model.device)
      all_samples.extend(list(text_samples))
  generative_ppl = 0.
  entropy = 0.
  if not config.sampling.semi_ar:
    generative_ppl = model.metrics.gen_ppl.compute().item()
    entropy = model.metrics.sample_entropy.compute().item()
    logger.info(f'Generative perplexity: {generative_ppl}')
    logger.info(f'Sample entropy: {entropy}')
  samples_path = config.eval.generated_samples_path
  with fsspec.open(samples_path, 'w') as f:
    json.dump({'generative_ppl': generative_ppl,
               'entropy': entropy,
               'generated_seqs': all_samples}, f, indent=4)
  logger.info(f'Samples saved at: {samples_path}',)

def _eval_ppl(diffusion_model, config, logger, tokenizer):
  logger.info('Starting Perplexity Eval.')

  model = _load_from_checkpoint(
    diffusion_model=diffusion_model,
    config=config,
    tokenizer=tokenizer)
  if config.eval.disable_ema:
    logger.info('Disabling EMA.')
    model.ema = None

  wandb_logger = None
  if config.get('wandb', None) is not None:
    wandb_logger = L.pytorch.loggers.WandbLogger(
      config=omegaconf.OmegaConf.to_object(config),
      ** config.wandb)
  callbacks = []
  if 'callbacks' in config:
    for _, callback in config.callbacks.items():
      callbacks.append(hydra.utils.instantiate(callback))
  trainer = hydra.utils.instantiate(
    config.trainer,
    default_root_dir=os.getcwd(),
    callbacks=callbacks,
    strategy=hydra.utils.instantiate(config.strategy),
    logger=wandb_logger)
  _, valid_ds = dataloader.get_dataloaders(
    config, tokenizer, skip_train=True, valid_seed=config.seed)
  trainer.validate(model, valid_ds)


def _train(diffusion_model, config, logger, tokenizer):
  logger.info('Starting Training.')
  wandb_logger = None
  if config.get('wandb', None) is not None:
    wandb_logger = L.pytorch.loggers.WandbLogger(
      config=omegaconf.OmegaConf.to_object(config),
      **config.wandb)

  if (config.checkpointing.resume_from_ckpt
      and config.checkpointing.resume_ckpt_path is not None
      and utils.fsspec_exists(
        config.checkpointing.resume_ckpt_path)):
    ckpt_path = config.checkpointing.resume_ckpt_path
  else:
    ckpt_path = None

  # Lightning callbacks
  callbacks = []
  if 'callbacks' in config:
    for _, callback in config.callbacks.items():
      callbacks.append(hydra.utils.instantiate(callback))

  train_ds, valid_ds = dataloader.get_dataloaders(
    config, tokenizer)
  _print_batch(config, train_ds, valid_ds, tokenizer)

  if config.training.finetune_path != '':
    assert utils.fsspec_exists(config.training.finetune_path)
    model = diffusion_model.load_from_checkpoint(
      config.training.finetune_path,
      tokenizer=tokenizer,
      config=config)
  else:
    model = diffusion_model(config, tokenizer=valid_ds.tokenizer)

  trainer = hydra.utils.instantiate(
    config.trainer,
    default_root_dir=os.getcwd(),
    callbacks=callbacks,
    strategy=hydra.utils.instantiate(config.strategy),
    logger=wandb_logger)
  trainer.fit(model, train_ds, valid_ds, ckpt_path=ckpt_path)


def _eval_fid(diffusion_model, config, logger, tokenizer):
  logger.info('Preparing data and model for FID eval.')
  fabric = Fabric(accelerator=config.trainer.accelerator,
                  devices=config.trainer.devices,
                  num_nodes=config.trainer.num_nodes)
  
  fabric.launch()
  seed = config.seed + fabric.global_rank
  L.seed_everything(seed)
  print(f'(Rank {fabric.global_rank}): seed: {seed}')
  model = _load_from_checkpoint(
    diffusion_model=diffusion_model,
    config=config,
    tokenizer=tokenizer)
  model.to(fabric.device)
  if config.eval.disable_ema:
    logger.info('Disabling EMA.')
    model.ema = None
  model._eval_mode()

  assert config.data.train == 'cifar10', \
                    'FID eval only implemented for CIFAR-10'

  # Like in flow matching papers: FID against train
  loader, _ = dataloader.get_dataloaders(config, 
    tokenizer=tokenizer, skip_valid=True)

  sampler = DistributedSampler(
    loader.dataset,
    num_replicas=fabric.world_size,
    rank=fabric.global_rank,
    shuffle=False)
  
  loader = torch.utils.data.DataLoader(
    loader.dataset,
    batch_size=config.loader.eval_batch_size,
    sampler=sampler,
    num_workers=loader.num_workers if hasattr(loader, 'num_workers') else 0,
    pin_memory=getattr(loader, 'pin_memory', False))
  
  # Check each GPU must generate the same number of images
  assert len(loader) == len(loader.dataset) // loader.batch_size // fabric.world_size, \
     f'{len(loader)=}, {len(loader.dataset)=}, {loader.batch_size=}, {fabric.world_size=}'

  fid_calculator = FrechetInceptionDistance(
    normalize=False).to(fabric.device)
  is_calculator = InceptionScore(
    normalize=False).to(fabric.device)
  
  desc = f'(Rank {fabric.global_rank}) Sampling...'
  for batch in tqdm(loader, desc=desc):
    real_samples = batch['input_ids']
    # Generate images with labels matching the true data
    labels = batch['labels']

    gen_samples = model.generate_samples(
      num_samples=real_samples.shape[0], 
      num_steps=config.sampling.steps,
      labels=labels)
    # Reshape 1D seq -> 2D image
    gen_samples = model.tokenizer.batch_decode(gen_samples)
    real_samples = model.tokenizer.batch_decode(
      real_samples).to(fabric.device)

    fid_calculator.update(gen_samples, real=False)
    fid_calculator.update(real_samples, real=True)
    is_calculator.update(gen_samples)
  
  fabric.barrier()
  logger.info('Done sampling. Computing FID & IS...')
  fid = fid_calculator.compute()
  incep_score = is_calculator.compute()

  if fabric.global_rank == 0:
    logger.info(f'FID: {fid}')
    logger.info(f'IS: {incep_score}')
  fabric.barrier()


@hydra.main(version_base=None, config_path='configs',
            config_name='config')
def main(config):
  """Main entry point for training."""
  L.seed_everything(config.seed)
  _print_config(config, resolve=True, save_cfg=True)
  
  logger = utils.get_logger(__name__)
  tokenizer = dataloader.get_tokenizer(config)
  if config.algo.name == 'ar':
    diffusion_model = algo.AR
  elif config.algo.name == 'mdlm':
    diffusion_model = algo.MDLM
  elif config.algo.name == 'duo_base':
    diffusion_model = algo.DUO_BASE
  elif config.algo.name == 'd3pm':
    diffusion_model = algo.D3PMAbsorb
  elif config.algo.name == 'sedd':
    diffusion_model = algo.SEDDAbsorb
  elif config.algo.name == 'duo':
    diffusion_model = algo.DUO
  elif config.algo.name == 'distillation':
    diffusion_model = algo.Distillation
  elif config.algo.name == 'ot-finetune':
    diffusion_model = algo.OptimalTransportFinetune
  else:
    raise ValueError(
      f'Invalid algorithm name: {config.algo.name}')
  kwargs = {'diffusion_model': diffusion_model,
            'config': config,
            'tokenizer': tokenizer,
            'logger': logger}
  if config.mode == 'sample_eval':
    _generate_samples(**kwargs)
  elif config.mode == 'ppl_eval':
    _eval_ppl(**kwargs)
  elif config.mode == 'fid_eval':
    _eval_fid(**kwargs)
  else:
    _train(**kwargs)


if __name__ == '__main__':
  main()

	import json
	import os

	import fsspec
	import hydra
	import lightning as L
	from lightning.fabric import Fabric
	import omegaconf
	import rich.syntax
	import rich.tree
	import torch
	from torch.utils.data.distributed import DistributedSampler
	from torchmetrics.image.fid import FrechetInceptionDistance
	from torchmetrics.image.inception import InceptionScore
	from tqdm import tqdm, trange

	import algo
	import dataloader
	import utils

	omegaconf.OmegaConf.register_new_resolver(
	'cwd', os.getcwd)
	omegaconf.OmegaConf.register_new_resolver(
	'device_count', torch.cuda.device_count)
	omegaconf.OmegaConf.register_new_resolver(
	'eval', eval)
	omegaconf.OmegaConf.register_new_resolver(
	'div_up', lambda x, y: (x + y - 1) // y)


	def _load_from_checkpoint(diffusion_model, config, tokenizer):
	if 'hf' in config.algo.backbone:
	return diffusion_model(
	config, tokenizer=tokenizer).to('cuda')

	return diffusion_model.load_from_checkpoint(
	config.eval.checkpoint_path,
	tokenizer=tokenizer,
	config=config)


	@L.pytorch.utilities.rank_zero_only
	def _print_config(
	config: omegaconf.DictConfig,
	resolve: bool = True,
	save_cfg: bool = True) -> None:
	"""Prints content of DictConfig using Rich library and its tree structure.

	Args:
	config (DictConfig): Configuration composed by Hydra.
	resolve (bool): Whether to resolve reference fields of DictConfig.
	save_cfg (bool): Whether to save the configuration tree to a file.
	"""

	style = 'dim'
	tree = rich.tree.Tree('CONFIG', style=style, guide_style=style)

	fields = config.keys()
	for field in fields:
	branch = tree.add(field, style=style, guide_style=style)

	config_section = config.get(field)
	branch_content = str(config_section)
	if isinstance(config_section, omegaconf.DictConfig):
	branch_content = omegaconf.OmegaConf.to_yaml(
	config_section, resolve=resolve)

	branch.add(rich.syntax.Syntax(branch_content, 'yaml'))
	rich.print(tree)
	if save_cfg:
	with fsspec.open(
	'{}/config_tree.txt'.format(
	config.checkpointing.save_dir), 'w') as fp:
	rich.print(tree, file=fp)


	@L.pytorch.utilities.rank_zero_only
	def _print_batch(config, train_ds, valid_ds, tokenizer, k=64):
	for dl_type, dl in [
	('train', train_ds), ('valid', valid_ds)]:
	print(f'Printing {dl_type} dataloader batch.')
	batch = next(iter(dl))
	print('Batch input_ids.shape', batch['input_ids'].shape)
	if config.data.modality == 'text':
	first = batch['input_ids'][0, :k]
	last = batch['input_ids'][0, -k:]
	print(f'First {k} tokens:', tokenizer.decode(first))
	print('ids:', first)
	print(f'Last {k} tokens:', tokenizer.decode(last))
	print('ids:', last)


	def _generate_samples(diffusion_model, config, logger,
	tokenizer):
	logger.info('Starting Sample Eval.')
	model = _load_from_checkpoint(
	diffusion_model=diffusion_model,
	config=config,
	tokenizer=tokenizer)
	model.metrics.gen_ppl.reset()
	model.metrics.sample_entropy.reset()
	if config.eval.disable_ema:
	logger.info('Disabling EMA.')
	model.ema = None
	stride_length = config.sampling.stride_length
	num_strides = config.sampling.num_strides
	all_samples = []
	for _ in trange(config.sampling.num_sample_batches):
	if config.sampling.semi_ar:
	_, intermediate_samples, _ = model.restore_model_and_semi_ar_sample(
	stride_length=stride_length,
	num_strides=num_strides,
	dt=1 / config.sampling.steps)
	text_samples = intermediate_samples[-1]
	# Note: Samples generated using semi-ar method
	# need to to be processed before computing generative perplexity
	# since these samples contain numerous <\|endoftext\|> tokens
	# and diffusion.compute_generative_perplexity() discards
	# any text after the first EOS token.
	else:
	samples = model.restore_model_and_sample(
	num_steps=config.sampling.steps)
	model.metrics.record_entropy(samples)
	text_samples = model.tokenizer.batch_decode(samples)
	model.metrics.record_generative_perplexity(
	text_samples, config.model.length, model.device)
	all_samples.extend(list(text_samples))
	generative_ppl = 0.
	entropy = 0.
	if not config.sampling.semi_ar:
	generative_ppl = model.metrics.gen_ppl.compute().item()
	entropy = model.metrics.sample_entropy.compute().item()
	logger.info(f'Generative perplexity: {generative_ppl}')
	logger.info(f'Sample entropy: {entropy}')
	samples_path = config.eval.generated_samples_path
	with fsspec.open(samples_path, 'w') as f:
	json.dump({'generative_ppl': generative_ppl,
	'entropy': entropy,
	'generated_seqs': all_samples}, f, indent=4)
	logger.info(f'Samples saved at: {samples_path}',)

	def _eval_ppl(diffusion_model, config, logger, tokenizer):
	logger.info('Starting Perplexity Eval.')

	model = _load_from_checkpoint(
	diffusion_model=diffusion_model,
	config=config,
	tokenizer=tokenizer)
	if config.eval.disable_ema:
	logger.info('Disabling EMA.')
	model.ema = None

	wandb_logger = None
	if config.get('wandb', None) is not None:
	wandb_logger = L.pytorch.loggers.WandbLogger(
	config=omegaconf.OmegaConf.to_object(config),
	** config.wandb)
	callbacks = []
	if 'callbacks' in config:
	for _, callback in config.callbacks.items():
	callbacks.append(hydra.utils.instantiate(callback))
	trainer = hydra.utils.instantiate(
	config.trainer,
	default_root_dir=os.getcwd(),
	callbacks=callbacks,
	strategy=hydra.utils.instantiate(config.strategy),
	logger=wandb_logger)
	_, valid_ds = dataloader.get_dataloaders(
	config, tokenizer, skip_train=True, valid_seed=config.seed)
	trainer.validate(model, valid_ds)


	def _train(diffusion_model, config, logger, tokenizer):
	logger.info('Starting Training.')
	wandb_logger = None
	if config.get('wandb', None) is not None:
	wandb_logger = L.pytorch.loggers.WandbLogger(
	config=omegaconf.OmegaConf.to_object(config),
	**config.wandb)

	if (config.checkpointing.resume_from_ckpt
	and config.checkpointing.resume_ckpt_path is not None
	and utils.fsspec_exists(
	config.checkpointing.resume_ckpt_path)):
	ckpt_path = config.checkpointing.resume_ckpt_path
	else:
	ckpt_path = None

	# Lightning callbacks
	callbacks = []
	if 'callbacks' in config:
	for _, callback in config.callbacks.items():
	callbacks.append(hydra.utils.instantiate(callback))

	train_ds, valid_ds = dataloader.get_dataloaders(
	config, tokenizer)
	_print_batch(config, train_ds, valid_ds, tokenizer)

	if config.training.finetune_path != '':
	assert utils.fsspec_exists(config.training.finetune_path)
	model = diffusion_model.load_from_checkpoint(
	config.training.finetune_path,
	tokenizer=tokenizer,
	config=config)
	else:
	model = diffusion_model(config, tokenizer=valid_ds.tokenizer)

	trainer = hydra.utils.instantiate(
	config.trainer,
	default_root_dir=os.getcwd(),
	callbacks=callbacks,
	strategy=hydra.utils.instantiate(config.strategy),
	logger=wandb_logger)
	trainer.fit(model, train_ds, valid_ds, ckpt_path=ckpt_path)


	def _eval_fid(diffusion_model, config, logger, tokenizer):
	logger.info('Preparing data and model for FID eval.')
	fabric = Fabric(accelerator=config.trainer.accelerator,
	devices=config.trainer.devices,
	num_nodes=config.trainer.num_nodes)

	fabric.launch()
	seed = config.seed + fabric.global_rank
	L.seed_everything(seed)
	print(f'(Rank {fabric.global_rank}): seed: {seed}')
	model = _load_from_checkpoint(
	diffusion_model=diffusion_model,
	config=config,
	tokenizer=tokenizer)
	model.to(fabric.device)
	if config.eval.disable_ema:
	logger.info('Disabling EMA.')
	model.ema = None
	model._eval_mode()

	assert config.data.train == 'cifar10', \
	'FID eval only implemented for CIFAR-10'

	# Like in flow matching papers: FID against train
	loader, _ = dataloader.get_dataloaders(config,
	tokenizer=tokenizer, skip_valid=True)

	sampler = DistributedSampler(
	loader.dataset,
	num_replicas=fabric.world_size,
	rank=fabric.global_rank,
	shuffle=False)

	loader = torch.utils.data.DataLoader(
	loader.dataset,
	batch_size=config.loader.eval_batch_size,
	sampler=sampler,
	num_workers=loader.num_workers if hasattr(loader, 'num_workers') else 0,
	pin_memory=getattr(loader, 'pin_memory', False))

	# Check each GPU must generate the same number of images
	assert len(loader) == len(loader.dataset) // loader.batch_size // fabric.world_size, \
	f'{len(loader)=}, {len(loader.dataset)=}, {loader.batch_size=}, {fabric.world_size=}'

	fid_calculator = FrechetInceptionDistance(
	normalize=False).to(fabric.device)
	is_calculator = InceptionScore(
	normalize=False).to(fabric.device)

	desc = f'(Rank {fabric.global_rank}) Sampling...'
	for batch in tqdm(loader, desc=desc):
	real_samples = batch['input_ids']
	# Generate images with labels matching the true data
	labels = batch['labels']

	gen_samples = model.generate_samples(
	num_samples=real_samples.shape[0],
	num_steps=config.sampling.steps,
	labels=labels)
	# Reshape 1D seq -> 2D image
	gen_samples = model.tokenizer.batch_decode(gen_samples)
	real_samples = model.tokenizer.batch_decode(
	real_samples).to(fabric.device)

	fid_calculator.update(gen_samples, real=False)
	fid_calculator.update(real_samples, real=True)
	is_calculator.update(gen_samples)

	fabric.barrier()
	logger.info('Done sampling. Computing FID & IS...')
	fid = fid_calculator.compute()
	incep_score = is_calculator.compute()

	if fabric.global_rank == 0:
	logger.info(f'FID: {fid}')
	logger.info(f'IS: {incep_score}')
	fabric.barrier()


	@hydra.main(version_base=None, config_path='configs',
	config_name='config')
	def main(config):
	"""Main entry point for training."""
	L.seed_everything(config.seed)
	_print_config(config, resolve=True, save_cfg=True)

	logger = utils.get_logger(__name__)
	tokenizer = dataloader.get_tokenizer(config)
	if config.algo.name == 'ar':
	diffusion_model = algo.AR
	elif config.algo.name == 'mdlm':
	diffusion_model = algo.MDLM
	elif config.algo.name == 'duo_base':
	diffusion_model = algo.DUO_BASE
	elif config.algo.name == 'd3pm':
	diffusion_model = algo.D3PMAbsorb
	elif config.algo.name == 'sedd':
	diffusion_model = algo.SEDDAbsorb
	elif config.algo.name == 'duo':
	diffusion_model = algo.DUO
	elif config.algo.name == 'distillation':
	diffusion_model = algo.Distillation
	elif config.algo.name == 'ot-finetune':
	diffusion_model = algo.OptimalTransportFinetune
	else:
	raise ValueError(
	f'Invalid algorithm name: {config.algo.name}')
	kwargs = {'diffusion_model': diffusion_model,
	'config': config,
	'tokenizer': tokenizer,
	'logger': logger}
	if config.mode == 'sample_eval':
	_generate_samples(**kwargs)
	elif config.mode == 'ppl_eval':
	_eval_ppl(**kwargs)
	elif config.mode == 'fid_eval':
	_eval_fid(**kwargs)
	else:
	_train(**kwargs)


	if __name__ == '__main__':
	main()