"""Computes and stores the average and current value"""

def __init__(self):

self.reset()

def reset(self):

self.val = 0

self.avg = 0

self.sum = 0

self.count = 0

def update(self, val, n=1):

self.val = val

self.sum += val * n

self.count += n

return torch.autograd.grad(

outputs=loss,

inputs=layer,

grad_outputs=torch.ones_like(loss),

retain_graph=True,

gradients = torch.autograd.grad(

outputs=output,

inputs=images,

grad_outputs=torch.ones(output.size(), device=images.device),

create_graph=True,

retain_graph=True,

only_inputs=True,

)[0]

bsz = gradients.shape[0]

gradients = torch.reshape(gradients, (bsz, -1))

save_path = Path(args.output_dir) / f"checkpoint-{global_step}"

# retrieve the model on all processes for deepspeed stage 3 to work then save on one process (we are not using stage 3 yet)

# XXX: could also make this conditional on deepspeed

state_dict = accelerator.get_state_dict(model)

discr_state_dict = accelerator.get_state_dict(discriminator)

if accelerator.is_main_process:

unwrapped_model = accelerator.unwrap_model(model)

unwrapped_model.save_pretrained(

save_path / "unwrapped_model",

save_function=accelerator.save,

state_dict=state_dict,

)

torch.save(discr_state_dict, save_path / "unwrapped_discriminator")

json.dump({"global_step": global_step}, (save_path / "metadata.json").open("w+"))

logger.info(f"Saved state to {save_path}")

accelerator.save_state(save_path)

if args.latest_checkpoint_only:

latest_checkpoint_path = Path(args.output_dir) / f"checkpoint-{global_step-args.checkpointing_steps}"

if accelerator.is_main_process:

if latest_checkpoint_path.exists():

for name, param in model.named_parameters():

if param.grad is not None:

grads = param.grad.detach().data

grad_norm = (grads.norm(p=2) / grads.numel()).item()

parser = argparse.ArgumentParser(description="Simple example of a training script.")

parser.add_argument("--log_grad_norm_steps", type=int, default=500,

help=("Print logs of gradient norms every X steps."))

parser.add_argument("--log_steps", type=int, default=50, help=("Print logs every X steps."))

parser.add_argument("--validation_steps", type=int, default=5000,

help=(

"Run validation every X steps. Validation consists of running reconstruction on images in"

" `args.validation_images` and logging the reconstructed images."

),

)

parser.add_argument("--log_image_steps", type=int, default=100)

parser.add_argument("--vae_loss", type=str, default="l1", help="The loss function for vae reconstruction loss.")

parser.add_argument("--pretrained_model_name_or_path", type=str, default=None,

help="Path to pretrained model or model identifier from huggingface.co/models.")

parser.add_argument("--model_config_name_or_path", type=str, default=None,

help="The config of the Vq model to train, leave as None to use standard DDPM configuration.")

parser.add_argument("--discriminator_config_name_or_path", type=str, default=None,

help="The config of the discriminator model to train, leave as None to use standard DDPM configuration.")

parser.add_argument("--dataset_name", type=str, default="robotic",

help=(

"The name of the Dataset (from the HuggingFace hub) to train on (could be your own, possibly private,"

" dataset). It can also be a path pointing to a local copy of a dataset in your filesystem,"

" or to a folder containing files that 🤗 Datasets can understand."

),

)

parser.add_argument("--output_dir", type=str, default="vqgan-output",

help="The output directory where the model predictions and checkpoints will be written.")

parser.add_argument("--cache_dir", type=str, default=None,

help="The directory where the downloaded models and datasets will be stored.")

parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")

parser.add_argument("--resolution", type=int, default=256,

help=(

"The resolution for input images, all the images in the train/validation dataset will be resized to this"

" resolution"

),

)

parser.add_argument("--train_batch_size", type=int, default=16,

help="Batch size (per device) for the training dataloader.")

parser.add_argument("--num_train_epochs", type=int, default=1)

parser.add_argument("--max_train_steps", type=int, default=1000000,

help="Total number of training steps to perform. If provided, overrides num_train_epochs.")

parser.add_argument("--gradient_accumulation_steps", type=int, default=1,

help="Number of updates steps to accumulate before performing a backward/update pass.")

parser.add_argument("--gradient_checkpointing", action="store_true",

help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.")

parser.add_argument("--discr_learning_rate", type=float, default=1e-4,

help="Initial learning rate (after the potential warmup period) to use.",)

parser.add_argument("--learning_rate", type=float, default=1e-4,

help="Initial learning rate (after the potential warmup period) to use.")

parser.add_argument("--scale_lr", action="store_true", default=False,

help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.")

parser.add_argument("--lr_scheduler", type=str, default="constant_with_warmup",

help=(

'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'

' "constant", "constant_with_warmup"]'

),

)

parser.add_argument("--discr_lr_scheduler", type=str, default="constant_with_warmup",

help=(

'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'

' "constant", "constant_with_warmup"]'

),

)

parser.add_argument("--lr_warmup_steps", type=int, default=500,

help="Number of steps for the warmup in the lr scheduler.")

parser.add_argument("--use_8bit_adam", action="store_true",

help="Whether or not to use 8-bit Adam from bitsandbytes.")

parser.add_argument("--allow_tf32", action="store_true",

help=(

"Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"

" https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"

),

)

parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.")

parser.add_argument("--dataloader_num_workers", type=int, default=0,

help=(

"Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process."

),

)

parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")

parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")

parser.add_argument("--adam_weight_decay", type=float, default=0, help="Weight decay to use.")

parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")

parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")

parser.add_argument("--logging_dir", type=str, default="logs",

help=(

"[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"

" *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."

),

)

parser.add_argument("--mixed_precision", type=str, default=None, choices=["no", "fp16", "bf16"],

help=(

"Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="

" 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the"

" flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."

),

)

parser.add_argument("--report_to", type=str, default="tensorboard",

help=(

'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'

' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'

),

)

parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")

parser.add_argument("--checkpointing_steps", type=int, default=5000,

help=(

"Save a checkpoint of the training state every X updates. These checkpoints are only suitable for resuming"

" training using `--resume_from_checkpoint`."

),

)

parser.add_argument("--resume_from_checkpoint", type=str, default=None,

help=(

"Whether training should be resumed from a previous checkpoint. Use a path saved by"

' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'

),

)

parser.add_argument("--enable_xformers_memory_efficient_attention",

action="store_true", help="Whether or not to use xformers.")

parser.add_argument("--tracker_project_name", type=str, default="vqgan-training",

help=(

"The `project_name` argument passed to Accelerator.init_trackers for"

" more information see https://huggingface.co/docs/accelerate/v0.17.0/en/package_reference/accelerator#accelerate.Accelerator"

),

)

parser.add_argument("--segment_length", type=int, default=5,

help="The length of the segmented trajectories to use for the training.")

parser.add_argument("--context_length", type=int, default=1)

parser.add_argument("--segment_horizon", type=int, default=16)

parser.add_argument('--video_stepsize', default=1, type=int)

parser.add_argument('--rand_select', default=False, action='store_true')

parser.add_argument('--rand_shuffle', default=False, action='store_true')

parser.add_argument('--model_type', default='vqgan', type=str, choices=['vqgan', 'ctx_vqgan'], help='Type of model to use')

parser.add_argument('--dataset_path', default='/data2/tensorflow_datasets',

type=str, help='Path to the tensorflow datasets')

parser.add_argument('--dataset_size', default=None, type=int)

parser.add_argument('--weighted_mse', default=None, type=float)

parser.add_argument('--weighted_gan', default=False, action='store_true')

parser.add_argument('--disc_start', default=0, type=int)

parser.add_argument('--disc_weight', default=0.8, type=float)

parser.add_argument('--latest_checkpoint_only', default=False, action='store_true')

parser.add_argument('--exp_name', default=None, type=str)

parser.add_argument('--disc_depth', default=4, type=int)

parser.add_argument('--perc_weight', default=1.0, type=float)

parser.add_argument('--recon_weight', default=1.0, type=float)

parser.add_argument('--oxe_data_mixes_type', default='frac', type=str)

parser.add_argument('--strong_aug', default=False, action='store_true')

parser.add_argument('--sthsth_root_path',

default='/data/something-something-v2/20bn-something-something-v2-frames-64', type=str)

parser.add_argument('--skip_first_val', default=False, action='store_true')

parser.add_argument('--start_global_step', default=0, type=int)

parser.add_argument('--balanced_loss', default=False, action='store_true')

parser.add_argument('--selected_params', default=False, action='store_true')

parser.add_argument('--no_aug', default=False, action='store_true')

args = parser.parse_args()

env_local_rank = int(os.environ.get("LOCAL_RANK", -1))

if env_local_rank != -1 and env_local_rank != args.local_rank:

args.local_rank = env_local_rank

#########################

# SETUP Accelerator #

#########################

args = parse_args()

args.output_dir = os.path.join(args.output_dir, time.strftime(

"%Y-%m-%d-%X", time.localtime()) + ("" if args.exp_name is None else f"-{args.exp_name}"))

os.makedirs(args.output_dir, exist_ok=True)

# Enable TF32 on Ampere GPUs

if args.allow_tf32:

torch.backends.cuda.matmul.allow_tf32 = True

torch.backends.cudnn.benchmark = True

torch.backends.cudnn.deterministic = False

logging_dir = os.path.join(args.output_dir, args.logging_dir)

os.makedirs(logging_dir, exist_ok=True)

accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)

accelerator = Accelerator(

gradient_accumulation_steps=args.gradient_accumulation_steps,

mixed_precision=args.mixed_precision,

log_with=args.report_to,

project_config=accelerator_project_config,

)

if accelerator.distributed_type == DistributedType.DEEPSPEED:

accelerator.state.deepspeed_plugin.deepspeed_config["train_micro_batch_size_per_gpu"] = args.train_batch_size

#####################################

# SETUP LOGGING, SEED and CONFIG #

#####################################

if accelerator.is_main_process:

tracker_config = dict(vars(args))

accelerator.init_trackers(args.tracker_project_name, tracker_config)

# If passed along, set the training seed now.

if args.seed is not None:

set_seed(args.seed, device_specific=True)

# Handle the repository creation

if accelerator.is_main_process:

if args.output_dir is not None:

os.makedirs(args.output_dir, exist_ok=True)

with open(os.path.join(args.output_dir, "cmd.sh"), "w") as f:

f.write("python " + " ".join(sys.argv))

src_path = os.path.join(args.output_dir, 'src')

os.makedirs(src_path, exist_ok=True)

os.system(f"rsync -rv --exclude-from=.gitignore . {src_path}")

#########################

# MODELS and OPTIMIZER #

#########################

logger.info("Loading models and optimizer")

if args.model_config_name_or_path is None and args.pretrained_model_name_or_path is None:

if args.model_type == "ctx_vqgan":

if args.resolution == 64:

config = json.load(open("configs/ctx_vae64/config.json"))

config.update({"context_length": args.context_length})

model = CompressiveVQModel(**config)

elif args.resolution == 256:

config = json.load(open("configs/ctx_vae/config.json"))

model = CompressiveVQModel(**config)

else:

raise NotImplementedError

else:

raise NotImplementedError

config = json.load(open("configs/vae/config.json"))

model = VQModel(**config)

elif args.pretrained_model_name_or_path is not None:

if args.model_type == "ctx_vqgan":

model = CompressiveVQModel.from_pretrained(

args.pretrained_model_name_or_path, subfolder=None, revision=None, variant=None, use_safetensor=True,

low_cpu_mem_usage=False, device_map=None,

ignore_mismatched_sizes=True

)

if args.pretrained_model_name_or_path == "configs/ctx_vae":

model.init_modules()

if args.context_length != model.context_length:

print(

f"[Warning] pretrained context length mismatch, change from {model.context_length} to {args.context_length}")

model.set_context_length(args.context_length)

elif args.model_type == "vqgan":

raise NotImplementedError

model = VQModel.from_pretrained(

args.pretrained_model_name_or_path, subfolder=None, revision=None, variant=None, use_safetensor=True

)

else:

raise NotImplementedError

else:

raise NotImplementedError

config = VQModel.load_config(args.model_config_name_or_path)

model = VQModel.from_config(config)

if args.use_ema:

ema_model = EMAModel(model.parameters(), model_cls=VQModel, model_config=model.config)

if args.discriminator_config_name_or_path is None:

discriminator = Discriminator(depth=args.disc_depth)

else:

discriminator = Discriminator(depth=args.disc_depth)

discriminator.load_state_dict(torch.load(args.discriminator_config_name_or_path))

# Perceptual loss

lpips = LPIPS().to(accelerator.device).eval()

# Enable flash attention if asked

if args.enable_xformers_memory_efficient_attention:

model.enable_xformers_memory_efficient_attention()

learning_rate = args.learning_rate

if args.scale_lr:

learning_rate = (

learning_rate * args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps

)

# Initialize the optimizer

if args.use_8bit_adam:

try:

import bitsandbytes as bnb

except ImportError:

raise ImportError(

"Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`"

)

optimizer_cls = bnb.optim.AdamW8bit

else:

optimizer_cls = torch.optim.AdamW

if args.selected_params:

# frozon codebook params

if args.model_type == 'ctx_vqgan':

params = [parameter for name, parameter in model.named_parameters() if 'quantize' not in name]

else:

raise NotImplementedError

else:

params = list(model.parameters())

optimizer = optimizer_cls(

params,

lr=args.learning_rate,

betas=(args.adam_beta1, args.adam_beta2),

weight_decay=args.adam_weight_decay,

eps=args.adam_epsilon,

)

discr_optimizer = optimizer_cls(

list(discriminator.parameters()),

lr=args.discr_learning_rate,

betas=(args.adam_beta1, args.adam_beta2),

weight_decay=args.adam_weight_decay,

eps=args.adam_epsilon,

)

##################################

# DATLOADER and LR-SCHEDULER #

#################################

logger.info("Creating dataloaders and lr_scheduler")

total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps

# DataLoaders creation:

if args.dataset_name != "robotic":

raise NotImplementedError

# DataLoaders creation:

if args.strong_aug:

augmentation_args = {

'brightness': [0.6, 1.4],

'contrast': [0.6, 1.4],

'saturation': [0.6, 1.4],

'hue': [-0.5, 0.5],

'random_resized_crop_scale': (0.6, 1.0),

'random_resized_crop_ratio': (0.75, 1.3333),

'no_aug': args.no_aug,

}

else:

augmentation_args = {

'brightness': [0.9, 1.1],

'contrast': [0.9, 1.1],

'saturation': [0.9, 1.1],

'hue': [-0.05, 0.05],

'random_resized_crop_scale': (0.8, 1.0),

'random_resized_crop_ratio': (0.9, 1.1),

'no_aug': args.no_aug,

}

segment_args = {

'random_selection': args.rand_select,

'random_shuffle': args.rand_shuffle,

'goal_conditioned': False,

'segment_length': args.segment_length,

'context_length': args.context_length,

'stepsize': args.video_stepsize,

'segment_horizon': args.segment_horizon,

}

train_dataloader = SimpleRoboticDataLoaderv2(

parent_dir=args.dataset_path,

datasets=DATASET_NAMED_MIXES[args.oxe_data_mixes_type],

batch_size=args.train_batch_size,

num_workers=args.dataloader_num_workers,

train=True,

maxsize=args.dataset_size,

image_size=args.resolution,

sthsth_root_path=args.sthsth_root_path,

**augmentation_args,

**segment_args,

)

eval_dataloader = SimpleRoboticDataLoaderv2(

parent_dir=args.dataset_path,

datasets=DATASET_NAMED_MIXES[args.oxe_data_mixes_type],

batch_size=args.train_batch_size,

num_workers=args.dataloader_num_workers,

train=False,

image_size=args.resolution,

sthsth_root_path=args.sthsth_root_path,

**augmentation_args,

**segment_args,

)

lr_scheduler = get_scheduler(

args.lr_scheduler,

optimizer=optimizer,

num_training_steps=args.max_train_steps * accelerator.num_processes,

num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,

)

discr_lr_scheduler = get_scheduler(

args.discr_lr_scheduler,

optimizer=discr_optimizer,

num_training_steps=args.max_train_steps * accelerator.num_processes,

num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,

)

# Prepare everything with accelerator

logger.info("Preparing model, optimizer and dataloaders")

# The dataloader are already aware of distributed training, so we don't need to prepare them.

model, discriminator, optimizer, discr_optimizer, lr_scheduler, discr_lr_scheduler = accelerator.prepare(

model, discriminator, optimizer, discr_optimizer, lr_scheduler, discr_lr_scheduler

)

# Train!

logger.info("***** Running training *****")

# logger.info(f" Num examples = {len(train_dataset)}")

logger.info(f" Num Epochs = {args.num_train_epochs}")

logger.info(f" Instantaneous batch size per device = {args.train_batch_size}")

logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")

logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}")

logger.info(f" Total optimization steps = {args.max_train_steps}")

global_step = args.start_global_step

first_epoch = 0

# Potentially load in the weights and states from a previous save

resume_from_checkpoint = args.resume_from_checkpoint

if resume_from_checkpoint:

if resume_from_checkpoint != "latest":

path = resume_from_checkpoint

else:

# Get the most recent checkpoint

dirs = os.listdir(args.output_dir)

dirs = [d for d in dirs if d.startswith("checkpoint")]

dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))

path = dirs[-1] if len(dirs) > 0 else None

path = os.path.join(args.output_dir, path)

if path is None:

accelerator.print(f"Checkpoint '{resume_from_checkpoint}' does not exist. Starting a new training run.")

resume_from_checkpoint = None

else:

accelerator.print(f"Resuming from checkpoint {path}")

accelerator.load_state(path)

accelerator.wait_for_everyone()

global_step = int(os.path.basename(path).split("-")[1])

# first_epoch = global_step // num_update_steps_per_epoch

batch_time_m = AverageMeter()

data_time_m = AverageMeter()

end = time.time()

progress_bar = tqdm(

range(0, args.max_train_steps),

initial=global_step,

desc="Steps",

# Only show the progress bar once on each machine.

disable=not accelerator.is_local_main_process,

)

# As stated above, we are not doing epoch based training here, but just using this for book keeping and being able to

# reuse the same training loop with other datasets/loaders.

avg_gen_loss, avg_discr_loss = None, None

avg_recon_loss, avg_commit_loss, avg_dyna_commit_loss, avg_perceptual_loss, avg_gan_loss, adaptive_weight, avg_residual_loss, avg_flow_loss = None, None, None, None, None, None, None, None

avg_ref_recon_loss, avg_ref_perceptual_loss = None, None

avg_feat_loss = None

avg_fake_logits, avg_real_logits = None, None

for epoch in range(first_epoch, args.num_train_epochs):

model.train()

for i, batch in enumerate(train_dataloader):

pixel_values = batch.to(accelerator.device, non_blocking=True)

pixel_values = pixel_values.reshape(-1, *pixel_values.shape[-3:])

original_pixel_values = pixel_values

generator_step = ((i // args.gradient_accumulation_steps) % 2) == 0

if generator_step:

data_time_m.update(time.time() - end)

# Train Step

# The behavior of accelerator.accumulate is to

# 1. Check if gradients are synced(reached gradient-accumulation_steps)

# 2. If so sync gradients by stopping the not syncing process

if generator_step:

optimizer.zero_grad(set_to_none=True)

else:

discr_optimizer.zero_grad(set_to_none=True)

# encode images to the latent space and get the commit loss from vq tokenization

# Return commit loss

if generator_step or global_step >= args.disc_start:

if 'ctx' in args.model_type:

with torch.no_grad():

BT, C, H, W = pixel_values.shape

B, T = (BT // args.segment_length), args.segment_length

frame_pixel_values = pixel_values.reshape(

args.train_batch_size, args.segment_length, C, H, W) # B, T, C, H, W

target = frame_pixel_values[:, args.context_length:].reshape(

B * (T - args.context_length), C, H, W) # B*(T-t), C, H, W

if args.context_length > 1:

reference_single = frame_pixel_values[:,

:args.context_length].reshape(-1, C, H, W) # B*t, C, H, W

reference = None # to raise not implementation error

else:

reference = frame_pixel_values[:, args.context_length - 1:args.context_length].repeat(

1, args.segment_length - args.context_length, 1, 1, 1).reshape(B * (T - args.context_length), C, H, W) # B*(T-t), C, H, W

reference_single = frame_pixel_values[:, args.context_length - 1]

pixel_values = target

if args.model_type == 'ctx_vqgan':

fmap, fmap_ref, commit_loss, dyna_commit_loss = model(sample=reference_single,

dyn_sample=target,

return_dict=False,

return_loss=True,

segment_len=args.segment_length - args.context_length)

else:

raise NotImplementedError

else:

fmap, commit_loss = model(pixel_values, return_dict=False, return_loss=True)

# weights for weighted losses

weights = None

weights_single = None

if generator_step:

with accelerator.accumulate(model):

def avg_loss(loss):

return accelerator.gather(loss.repeat(args.train_batch_size)).float().mean()

# reconstruction loss. Pixel level differences between input vs output

def get_recon_loss(gt, recon, weights):

if args.vae_loss == "l2":

loss = F.mse_loss(gt, recon, reduction='none')

else:

loss = F.l1_loss(gt, recon, reduction='none')

if weights is not None:

resized_weights = F.interpolate(weights, loss.shape[2:])

loss = (loss * resized_weights).mean()

else:

loss = loss.mean()

return loss

recon_loss = get_recon_loss(pixel_values, fmap, weights)

if args.balanced_loss:

loss = args.recon_weight * recon_loss * \

(args.segment_length - args.context_length) / args.segment_length

else:

loss = args.recon_weight * recon_loss

avg_recon_loss = avg_loss(recon_loss)

if 'ctx' in args.model_type:

ref_recon_loss = get_recon_loss(reference_single, fmap_ref, weights_single)

if args.balanced_loss:

loss += args.recon_weight * ref_recon_loss * args.context_length / args.segment_length

else:

loss += args.recon_weight * ref_recon_loss

avg_ref_recon_loss = avg_loss(ref_recon_loss)

# perceptual loss. The high level feature mean squared error loss

perceptual_loss = lpips(

pixel_values.contiguous() * 2 - 1.0,

fmap.contiguous() * 2 - 1.0,

weight=weights

).mean()

if args.balanced_loss:

loss += args.perc_weight * perceptual_loss * \

(args.segment_length - args.context_length) / args.segment_length

else:

loss += args.perc_weight * perceptual_loss

avg_perceptual_loss = avg_loss(perceptual_loss)

if 'ctx' in args.model_type:

ref_perceptual_loss = lpips(

reference_single.contiguous() * 2 - 1.0,

fmap_ref.contiguous() * 2 - 1.0,

weight=weights

).mean()

if args.balanced_loss:

loss += args.perc_weight * ref_perceptual_loss * args.context_length / args.segment_length

else:

loss += args.perc_weight * ref_perceptual_loss

avg_ref_perceptual_loss = avg_loss(ref_perceptual_loss)

# generator loss

if global_step >= args.disc_start:

if 'ctx' in args.model_type:

disc_fmap = torch.cat([fmap_ref, fmap], dim=0)

if weights_single is None or weights is None:

disc_weights = None

else:

disc_weights = torch.cat([weights_single, weights], dim=0)

else:

disc_fmap = fmap

disc_weights = weights

if disc_weights is not None and args.weighted_gan:

logits = discriminator(disc_fmap)

resized_weights = F.interpolate(disc_weights, logits.shape[2:])

gen_loss = -(resized_weights * logits).mean()

else:

gen_loss = -discriminator(disc_fmap).mean()

last_dec_layer = accelerator.unwrap_model(model).cond_decoder.conv_out.weight

norm_grad_wrt_perceptual_loss = grad_layer_wrt_loss(perceptual_loss, last_dec_layer).norm(p=2)

norm_grad_wrt_gen_loss = grad_layer_wrt_loss(gen_loss, last_dec_layer).norm(p=2)

adaptive_weight = norm_grad_wrt_perceptual_loss / norm_grad_wrt_gen_loss.clamp(min=1e-8)

adaptive_weight = adaptive_weight.clamp(max=1e4)

loss += args.disc_weight * adaptive_weight * gen_loss

avg_gan_loss = avg_loss(gen_loss)

# regularization losses

loss += commit_loss

avg_commit_loss = avg_loss(commit_loss)

if 'ctx' in args.model_type:

loss += dyna_commit_loss

avg_dyna_commit_loss = avg_loss(dyna_commit_loss)

# Gather thexd losses across all processes for logging (if we use distributed training).

avg_gen_loss = avg_loss(loss)

accelerator.backward(loss)

# print("detect unused_parameters for debug")

# for name, param in model.named_parameters():

# if param.grad is None:

# print(name)

if args.max_grad_norm is not None and accelerator.sync_gradients:

accelerator.clip_grad_norm_(model.parameters(), args.max_grad_norm)

optimizer.step()

lr_scheduler.step()

# log gradient norm before zeroing it

if accelerator.sync_gradients and global_step % args.log_grad_norm_steps == 0 and accelerator.is_main_process:

log_grad_norm(model, accelerator, global_step)

else:

# Return discriminator loss

with accelerator.accumulate(discriminator):

if global_step >= args.disc_start:

fmap.detach_()

if 'ctx' in args.model_type:

fmap_ref.detach_()

# pixel_values.requires_grad_()

if 'ctx' in args.model_type:

disc_pixel_values = torch.cat([reference_single, pixel_values], dim=0)

disc_fmap = torch.cat([fmap_ref, fmap], dim=0)

if weights_single is None or weights is None:

disc_weights = None

else:

disc_weights = torch.cat([weights_single, weights], dim=0)

else:

disc_pixel_values = pixel_values

disc_fmap = fmap

disc_weights = weights

real = discriminator(disc_pixel_values)

fake = discriminator(disc_fmap)

if weights is not None and args.weighted_gan:

resized_weights = F.interpolate(disc_weights, fake.shape[2:])

loss = (resized_weights * F.relu(1 + fake) + resized_weights * F.relu(1 - real)).mean()

else:

loss = (F.relu(1 + fake) + F.relu(1 - real)).mean()

# gp = gradient_penalty(pixel_values, real)

# loss += gp

if global_step < args.disc_start:

loss = loss * 0.0

avg_discr_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean()

if weights is not None and args.weighted_gan:

avg_fake_logits = accelerator.gather(

(resized_weights * fake).mean().repeat(args.train_batch_size)).mean()

avg_real_logits = accelerator.gather(

(resized_weights * real).mean().repeat(args.train_batch_size)).mean()

else:

avg_fake_logits = accelerator.gather(fake.mean().repeat(args.train_batch_size)).mean()

avg_real_logits = accelerator.gather(real.mean().repeat(args.train_batch_size)).mean()

accelerator.backward(loss)

if args.max_grad_norm is not None and accelerator.sync_gradients:

accelerator.clip_grad_norm_(discriminator.parameters(), args.max_grad_norm)

discr_optimizer.step()

discr_lr_scheduler.step()

if accelerator.sync_gradients and global_step % args.log_grad_norm_steps == 0 and accelerator.is_main_process:

log_grad_norm(discriminator, accelerator, global_step)

else:

pass # skip discriminator step if not started

# Checks if the accelerator has performed an optimization step behind the scenes

if accelerator.sync_gradients:

global_step += 1

progress_bar.update(1)

if args.use_ema:

ema_model.step(model.parameters())

if accelerator.sync_gradients and not generator_step and accelerator.is_main_process:

# wait for both generator and discriminator to settle

batch_time_m.update(time.time() - end)

progress_bar.set_postfix(batch_time=batch_time_m.val, data_time=data_time_m.val)

end = time.time()

# Log metrics

if global_step % args.log_steps == 0:

samples_per_second_per_gpu = (

args.gradient_accumulation_steps * args.train_batch_size / batch_time_m.val

)

logs = {

"step_discr_loss": avg_discr_loss.item() if avg_discr_loss is not None else 0.0,

"lr": lr_scheduler.get_last_lr()[0],

"samples/sec/gpu": samples_per_second_per_gpu,

"data_time": data_time_m.val,

"batch_time": batch_time_m.val,

}

if avg_gen_loss is not None:

logs["step_gen_loss"] = avg_gen_loss.item()

if avg_recon_loss is not None:

logs["gen_loss/step_recon_loss"] = avg_recon_loss.item()

if avg_ref_recon_loss is not None:

logs["gen_loss/step_ref_recon_loss"] = avg_ref_recon_loss.item()

if avg_commit_loss is not None:

logs["gen_loss/step_commit_loss"] = avg_commit_loss.item()

if avg_dyna_commit_loss is not None:

logs["gen_loss/step_dyna_commit_loss"] = avg_dyna_commit_loss.item()

if avg_perceptual_loss is not None:

logs["gen_loss/step_perceptual_loss"] = avg_perceptual_loss.item()

if avg_ref_perceptual_loss is not None:

logs["gen_loss/step_ref_perceptual_loss"] = avg_ref_perceptual_loss.item()

if avg_gan_loss is not None:

logs["gen_loss/step_gan_loss"] = avg_gan_loss.item()

if adaptive_weight is not None:

logs["gen_loss/adaptive_weight"] = adaptive_weight.item()

if avg_fake_logits is not None:

logs["disc_loss/step_fake_logits"] = avg_fake_logits.item()

if avg_real_logits is not None:

logs["disc_loss/step_real_logits"] = avg_real_logits.item()

if avg_fake_logits is not None and avg_real_logits is not None:

logs["disc_loss/step_logit_diff"] = avg_real_logits.item() - avg_fake_logits.item()

if avg_residual_loss is not None:

logs["gen_loss/step_residual_loss"] = avg_residual_loss.item()

if avg_flow_loss is not None:

logs["gen_loss/step_flow_loss"] = avg_flow_loss.item()

if avg_feat_loss is not None:

logs["gen_loss/step_feat_loss"] = avg_feat_loss.item()

logs["mem_used"] = psutil.virtual_memory().used / 1024 / 1024 / 1024

accelerator.log(logs, step=global_step)

# resetting batch / data time meters per log window

batch_time_m.reset()

data_time_m.reset()

# Save model checkpoint

if global_step % args.checkpointing_steps == 0:

save_checkpoint(model, discriminator, args, accelerator, global_step)

if accelerator.sync_gradients and generator_step and accelerator.is_main_process:

# Generate images

if global_step % args.log_image_steps == 1:

with torch.no_grad():

save_path = os.path.join(args.output_dir, "images", f"train-samples-{global_step}")

os.makedirs(save_path, exist_ok=True)

segment_length = args.segment_length - args.context_length

np_img = lambda x: x.detach().cpu().numpy().transpose(1, 2, 0)[:, :, ::-1] * 255

gt = np.concatenate([np_img(pixel_values[i]) for i in range(segment_length)], 1)

recon = np.concatenate([np_img(fmap[i]) for i in range(segment_length)], 1)

if 'ctx' in args.model_type:

if args.context_length > 1:

# fmap_ref: B, t, C, H, W -> B, C, H, t*W

ref_recon = fmap_ref.reshape(B, args.context_length, C, H, W)[0].permute(

1, 2, 0, 3).reshape(C, H, args.context_length * W)

ref_recon = np_img(ref_recon)

else:

ref_recon = np_img(fmap_ref[0])

diff = np.concatenate([np_img(fmap[i] - fmap[max(i - 1, 0)]) for i in range(segment_length)], 1)

# error = np.concatenate([np_img(torch.abs(fmap[i] - pixel_values[i])) for i in range(segment_length)], 1)

context = np.concatenate([np_img(original_pixel_values[i])

for i in range(args.context_length)], 1)

gt = np.concatenate([context, gt], 1)

recon = np.concatenate([np.zeros_like(context), recon], 1)

if 'ctx' in args.model_type:

if args.context_length > 1:

recon[:, :args.context_length * args.resolution, :] = ref_recon

else:

recon[:, (args.context_length - 1) * args.resolution:args.context_length *

args.resolution, :] = ref_recon

diff = np.concatenate([np.zeros_like(context), diff], 1)

# error = np.concatenate([np.zeros_like(context), error], 1)

error = np.abs(recon - gt)

cv2.imwrite(os.path.join(

save_path, f'train-samples-{global_step}.png'), np.concatenate([gt, recon, diff, error], 0))

# Validation

if global_step % args.validation_steps == 1 and (global_step > 1 or not args.skip_first_val):

with torch.no_grad():

model.eval()

recon_losses = []

perceptual_losses = []

val_iters = 100

bar = tqdm(range(val_iters), desc="validation")

for i, batch in enumerate(eval_dataloader):

if i == val_iters:

break

# preprocess

pixel_values = batch.to(accelerator.device, non_blocking=True)

pixel_values = pixel_values.reshape(-1, *pixel_values.shape[-3:])

original_pixel_values = pixel_values

BT, C, H, W = pixel_values.shape

B, T = (BT // args.segment_length), args.segment_length

frame_pixel_values = pixel_values.reshape(

args.train_batch_size, args.segment_length, C, H, W) # B, T, C, H, W

target = frame_pixel_values[:, args.context_length:].reshape(

B * (T - args.context_length), C, H, W) # B*(T-t), C, H, W

if args.context_length > 1:

# B*t, C, H, W

reference_single = frame_pixel_values[:, :args.context_length].reshape(-1, C, H, W)

reference = None # to raise not implementation error

else:

reference = frame_pixel_values[:, args.context_length - 1:args.context_length].repeat(