Spaces:
Running
Running
File size: 13,260 Bytes
94bafa8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 |
# All rights reserved.
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
"""
A minimal training script for Latte using PyTorch DDP.
"""
import torch
# Maybe use fp16 percision training need to set to False
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
import os
import math
import argparse
import torch.distributed as dist
from glob import glob
from time import time
from copy import deepcopy
from einops import rearrange
from models import get_models
from datasets import get_dataset
from models.clip import TextEmbedder
from diffusion import create_diffusion
from omegaconf import OmegaConf
from torch.utils.data import DataLoader
from diffusers.models import AutoencoderKL
from diffusers.optimization import get_scheduler
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
from utils import (clip_grad_norm_, create_logger, update_ema,
requires_grad, cleanup, create_tensorboard,
write_tensorboard, setup_distributed, get_experiment_dir)
#################################################################################
# Training Loop #
#################################################################################
def main(args):
assert torch.cuda.is_available(), "Training currently requires at least one GPU."
# Setup DDP:
setup_distributed()
rank = int(os.environ["RANK"])
local_rank = int(os.environ["LOCAL_RANK"])
device = torch.device("cuda", local_rank)
seed = args.global_seed + rank
torch.manual_seed(seed)
torch.cuda.set_device(device)
print(f"Starting rank={rank}, local rank={local_rank}, seed={seed}, world_size={dist.get_world_size()}.")
# Setup an experiment folder:
if rank == 0:
os.makedirs(args.results_dir, exist_ok=True) # Make results folder (holds all experiment subfolders)
experiment_index = len(glob(f"{args.results_dir}/*"))
model_string_name = args.model.replace("/", "-") # e.g., Latte-XL/2 --> Latte-XL-2 (for naming folders)
num_frame_string = 'F' + str(args.num_frames) + 'S' + str(args.frame_interval)
experiment_dir = f"{args.results_dir}/{experiment_index:03d}-{model_string_name}-{num_frame_string}-{args.dataset}" # Create an experiment folder
experiment_dir = get_experiment_dir(experiment_dir, args)
checkpoint_dir = f"{experiment_dir}/checkpoints" # Stores saved model checkpoints
os.makedirs(checkpoint_dir, exist_ok=True)
logger = create_logger(experiment_dir)
tb_writer = create_tensorboard(experiment_dir)
OmegaConf.save(args, os.path.join(experiment_dir, 'config.yaml'))
logger.info(f"Experiment directory created at {experiment_dir}")
else:
logger = create_logger(None)
tb_writer = None
# Create model:
assert args.image_size % 8 == 0, "Image size must be divisible by 8 (for the VAE encoder)."
sample_size = args.image_size // 8
args.latent_size = sample_size
model = get_models(args)
# Note that parameter initialization is done within the Latte constructor
ema = deepcopy(model).to(device) # Create an EMA of the model for use after training
requires_grad(ema, False)
diffusion = create_diffusion(timestep_respacing="") # default: 1000 steps, linear noise schedule
# vae = AutoencoderKL.from_pretrained(f"stabilityai/sd-vae-ft-ema").to(device)
vae = AutoencoderKL.from_pretrained(args.pretrained_model_path, subfolder="sd-vae-ft-mse").to(device)
# # use pretrained model?
if args.pretrained:
checkpoint = torch.load(args.pretrained, map_location=lambda storage, loc: storage)
if "ema" in checkpoint: # supports checkpoints from train.py
logger.info('Using ema ckpt!')
checkpoint = checkpoint["ema"]
model_dict = model.state_dict()
# 1. filter out unnecessary keys
pretrained_dict = {}
for k, v in checkpoint.items():
if k in model_dict:
pretrained_dict[k] = v
else:
logger.info('Ignoring: {}'.format(k))
logger.info('Successfully Load {}% original pretrained model weights '.format(len(pretrained_dict) / len(checkpoint.items()) * 100))
# 2. overwrite entries in the existing state dict
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
logger.info('Successfully load model at {}!'.format(args.pretrained))
if args.use_compile:
model = torch.compile(model)
if args.enable_xformers_memory_efficient_attention:
logger.info("Using Xformers!")
model.enable_xformers_memory_efficient_attention()
if args.gradient_checkpointing:
logger.info("Using gradient checkpointing!")
model.enable_gradient_checkpointing()
if args.fixed_spatial:
trainable_modules = (
"attn_temp",
)
model.requires_grad_(False)
for name, module in model.named_modules():
if name.endswith(tuple(trainable_modules)):
for params in module.parameters():
logger.info("WARNING: Only train {} parametes!".format(name))
params.requires_grad = True
logger.info("WARNING: Only train {} parametes!".format(trainable_modules))
# set distributed training
model = DDP(model.to(device), device_ids=[local_rank])
logger.info(f"Model Parameters: {sum(p.numel() for p in model.parameters()):,}")
opt = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0)
# Freeze vae and text_encoder
vae.requires_grad_(False)
# Setup data:
dataset = get_dataset(args)
sampler = DistributedSampler(
dataset,
num_replicas=dist.get_world_size(),
rank=rank,
shuffle=True,
seed=args.global_seed
)
loader = DataLoader(
dataset,
batch_size=int(args.local_batch_size),
shuffle=False,
sampler=sampler,
num_workers=args.num_workers,
pin_memory=True,
drop_last=True
)
logger.info(f"Dataset contains {len(dataset):,} videos ({args.webvideo_data_path})")
# Scheduler
lr_scheduler = get_scheduler(
name="constant",
optimizer=opt,
num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps,
num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
)
# Prepare models for training:
update_ema(ema, model.module, decay=0) # Ensure EMA is initialized with synced weights
model.train() # important! This enables embedding dropout for classifier-free guidance
ema.eval() # EMA model should always be in eval mode
# Variables for monitoring/logging purposes:
train_steps = 0
log_steps = 0
running_loss = 0
first_epoch = 0
start_time = time()
# We need to recalculate our total training steps as the size of the training dataloader may have changed.
num_update_steps_per_epoch = math.ceil(len(loader))
# Afterwards we recalculate our number of training epochs
num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
# Potentially load in the weights and states from a previous save
if args.resume_from_checkpoint:
# TODO, need to checkout
# Get the most recent checkpoint
dirs = os.listdir(os.path.join(experiment_dir, 'checkpoints'))
dirs = [d for d in dirs if d.endswith("pt")]
dirs = sorted(dirs, key=lambda x: int(x.split(".")[0]))
path = dirs[-1]
logger.info(f"Resuming from checkpoint {path}")
model.load_state(os.path.join(dirs, path))
train_steps = int(path.split(".")[0])
first_epoch = train_steps // num_update_steps_per_epoch
resume_step = train_steps % num_update_steps_per_epoch
for epoch in range(first_epoch, num_train_epochs):
sampler.set_epoch(epoch)
for step, video_data in enumerate(loader):
# Skip steps until we reach the resumed step
if args.resume_from_checkpoint and epoch == first_epoch and step < resume_step:
continue
x = video_data['video'].to(device, non_blocking=True)
video_name = video_data['video_name']
if args.dataset == "ucf101_img":
image_name = video_data['image_name']
image_names = []
for caption in image_name:
single_caption = [int(item) for item in caption.split('=====')]
image_names.append(torch.as_tensor(single_caption))
# x = x.to(device)
# y = y.to(device) # y is text prompt; no need put in gpu
with torch.no_grad():
# Map input images to latent space + normalize latents:
b, _, _, _, _ = x.shape
x = rearrange(x, 'b f c h w -> (b f) c h w').contiguous()
x = vae.encode(x).latent_dist.sample().mul_(0.18215)
x = rearrange(x, '(b f) c h w -> b f c h w', b=b).contiguous()
if args.extras == 78: # text-to-video
raise 'T2V training are Not supported at this moment!'
elif args.extras == 2:
if args.dataset == "ucf101_img":
model_kwargs = dict(y=video_name, y_image=image_names, use_image_num=args.use_image_num) # tav unet
else:
model_kwargs = dict(y=video_name) # tav unet
else:
model_kwargs = dict(y=None, use_image_num=args.use_image_num)
t = torch.randint(0, diffusion.num_timesteps, (x.shape[0],), device=device)
loss_dict = diffusion.training_losses(model, x, t, model_kwargs)
loss = loss_dict["loss"].mean()
loss.backward()
if train_steps < args.start_clip_iter: # if train_steps >= start_clip_iter, will clip gradient
gradient_norm = clip_grad_norm_(model.module.parameters(), args.clip_max_norm, clip_grad=False)
else:
gradient_norm = clip_grad_norm_(model.module.parameters(), args.clip_max_norm, clip_grad=True)
opt.step()
lr_scheduler.step()
opt.zero_grad()
update_ema(ema, model.module)
# Log loss values:
running_loss += loss.item()
log_steps += 1
train_steps += 1
if train_steps % args.log_every == 0:
# Measure training speed:
torch.cuda.synchronize()
end_time = time()
steps_per_sec = log_steps / (end_time - start_time)
# Reduce loss history over all processes:
avg_loss = torch.tensor(running_loss / log_steps, device=device)
dist.all_reduce(avg_loss, op=dist.ReduceOp.SUM)
avg_loss = avg_loss.item() / dist.get_world_size()
# logger.info(f"(step={train_steps:07d}) Train Loss: {avg_loss:.4f}, Train Steps/Sec: {steps_per_sec:.2f}")
logger.info(f"(step={train_steps:07d}/epoch={epoch:04d}) Train Loss: {avg_loss:.4f}, Gradient Norm: {gradient_norm:.4f}, Train Steps/Sec: {steps_per_sec:.2f}")
write_tensorboard(tb_writer, 'Train Loss', avg_loss, train_steps)
write_tensorboard(tb_writer, 'Gradient Norm', gradient_norm, train_steps)
# Reset monitoring variables:
running_loss = 0
log_steps = 0
start_time = time()
# Save Latte checkpoint:
if train_steps % args.ckpt_every == 0 and train_steps > 0:
if rank == 0:
checkpoint = {
# "model": model.module.state_dict(),
"ema": ema.state_dict(),
# "opt": opt.state_dict(),
# "args": args
}
checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}.pt"
torch.save(checkpoint, checkpoint_path)
logger.info(f"Saved checkpoint to {checkpoint_path}")
dist.barrier()
model.eval() # important! This disables randomized embedding dropout
# do any sampling/FID calculation/etc. with ema (or model) in eval mode ...
logger.info("Done!")
cleanup()
if __name__ == "__main__":
# Default args here will train Latte-XL/2 with the hyperparameters we used in our paper (except training iters).
parser = argparse.ArgumentParser()
parser.add_argument("--config", type=str, default="./configs/sky/sky_train.yaml")
args = parser.parse_args()
main(OmegaConf.load(args.config))
|