295 lines
12 KiB
Python
295 lines
12 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# 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 DiT using PyTorch DDP.
|
|
"""
|
|
import torch
|
|
# the first flag below was False when we tested this script but True makes A100 training a lot faster:
|
|
from torch import nn
|
|
|
|
torch.backends.cuda.matmul.allow_tf32 = True
|
|
torch.backends.cudnn.allow_tf32 = True
|
|
import torch.distributed as dist
|
|
from torch.nn.parallel import DistributedDataParallel as DDP
|
|
from torch.utils.data import DataLoader
|
|
from torch.utils.data.distributed import DistributedSampler
|
|
from collections import OrderedDict
|
|
from copy import deepcopy
|
|
from glob import glob
|
|
from time import time
|
|
import argparse
|
|
import logging
|
|
import os
|
|
from opensora.dataset import getdataset
|
|
from opensora.models.ae import getae
|
|
from opensora.utils.dataset_utils import Collate
|
|
from opensora.models.diffusion.dit.models import DiT_models as Diffusion_models
|
|
from opensora.models.diffusion.dit.diffusion import create_diffusion
|
|
|
|
#################################################################################
|
|
# Training Helper Functions #
|
|
#################################################################################
|
|
|
|
@torch.no_grad()
|
|
def update_ema(ema_model, model, decay=0.9999):
|
|
"""
|
|
Step the EMA model towards the current model.
|
|
"""
|
|
ema_params = OrderedDict(ema_model.named_parameters())
|
|
model_params = OrderedDict(model.named_parameters())
|
|
|
|
for name, param in model_params.items():
|
|
# TODO: Consider applying only to params that require_grad to avoid small numerical changes of pos_embed
|
|
ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)
|
|
|
|
|
|
def requires_grad(model, flag=True):
|
|
"""
|
|
Set requires_grad flag for all parameters in a model.
|
|
"""
|
|
for p in model.parameters():
|
|
p.requires_grad = flag
|
|
|
|
|
|
def cleanup():
|
|
"""
|
|
End DDP training.
|
|
"""
|
|
dist.destroy_process_group()
|
|
|
|
|
|
def create_logger(logging_dir):
|
|
"""
|
|
Create a logger that writes to a log file and stdout.
|
|
"""
|
|
if dist.get_rank() == 0: # real logger
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='[\033[34m%(asctime)s\033[0m] %(message)s',
|
|
datefmt='%Y-%m-%d %H:%M:%S',
|
|
handlers=[logging.StreamHandler(), logging.FileHandler(f"{logging_dir}/log.txt")]
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
else: # dummy logger (does nothing)
|
|
logger = logging.getLogger(__name__)
|
|
logger.addHandler(logging.NullHandler())
|
|
return logger
|
|
|
|
#################################################################################
|
|
# Training Loop #
|
|
#################################################################################
|
|
|
|
def main(args):
|
|
"""
|
|
Trains a new DiT model.
|
|
"""
|
|
assert torch.cuda.is_available(), "Training currently requires at least one GPU."
|
|
|
|
# Setup DDP:
|
|
dist.init_process_group("nccl")
|
|
assert args.global_batch_size % dist.get_world_size() == 0, f"Batch size must be divisible by world size."
|
|
rank = dist.get_rank()
|
|
device = rank % torch.cuda.device_count()
|
|
seed = args.global_seed * dist.get_world_size() + rank
|
|
torch.manual_seed(seed)
|
|
torch.cuda.set_device(device)
|
|
print(f"Starting rank={rank}, seed={seed}, world_size={dist.get_world_size()}.")
|
|
|
|
# Setup an experiment folder:
|
|
if rank == 0:
|
|
print(args)
|
|
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., DiT-XL/2 --> DiT-XL-2 (for naming folders)
|
|
experiment_dir = f"{args.results_dir}/{experiment_index:03d}-{model_string_name}" # Create an experiment folder
|
|
checkpoint_dir = f"{experiment_dir}/checkpoints" # Stores saved model checkpoints
|
|
os.makedirs(checkpoint_dir, exist_ok=True)
|
|
logger = create_logger(experiment_dir)
|
|
logger.info(f"Experiment directory created at {experiment_dir}")
|
|
else:
|
|
logger = create_logger(None)
|
|
|
|
# Create model:
|
|
vae_stride_t, vae_stride_h, vae_stride_w = [int(i) for i in args.ae[-5:].split('x')]
|
|
args.vae_stride_t, args.vae_stride_h, args.vae_stride_w = vae_stride_t, vae_stride_h, vae_stride_w
|
|
args.vae_stride = args.vae_stride_h
|
|
patch_size = args.model[-3:]
|
|
patch_size_t, patch_size_h, patch_size_w = int(patch_size[0]), int(patch_size[1]), int(patch_size[2])
|
|
args.patch_size = patch_size_h
|
|
args.patch_size_t, args.patch_size_h, args.patch_size_w = patch_size_t, patch_size_h, patch_size_w
|
|
assert args.num_frames % vae_stride_t == 0, "Image size must be divisible by 8 (for the VAE encoder)."
|
|
assert args.max_image_size % vae_stride_h == 0, "Image size must be divisible by 8 (for the VAE encoder)."
|
|
assert vae_stride_h == vae_stride_w, "Support now."
|
|
assert patch_size_h == patch_size_w, "Support now."
|
|
|
|
latent_size = (args.num_frames // vae_stride_t, args.max_image_size // vae_stride_h, args.max_image_size // vae_stride_w)
|
|
|
|
model = Diffusion_models[args.model](
|
|
input_size=latent_size,
|
|
num_classes=args.num_classes
|
|
)
|
|
model.gradient_checkpointing = args.gradient_checkpointing
|
|
|
|
|
|
if args.pt_ckpt is not None:
|
|
state_dict = torch.load(args.pt_ckpt, map_location='cpu')
|
|
if state_dict.get('model', None) is not None:
|
|
state_dict = state_dict['model']
|
|
del state_dict['x_embedder.proj.weight']
|
|
del state_dict['final_layer.linear.weight']
|
|
del state_dict['final_layer.linear.bias']
|
|
del state_dict['pos_embed']
|
|
del state_dict['y_embedder.embedding_table.weight']
|
|
missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
|
|
print('missing_keys:', missing_keys, 'unexpected_keys:', unexpected_keys)
|
|
assert len(missing_keys) == 4
|
|
logger.info(f"{missing_keys}, {unexpected_keys}")
|
|
|
|
# Note that parameter initialization is done within the DiT constructor
|
|
ema = deepcopy(model).to(device) # Create an EMA of the model for use after training
|
|
requires_grad(ema, False)
|
|
model = DDP(model.to(device), device_ids=[rank])
|
|
diffusion = create_diffusion(timestep_respacing="") # default: 1000 steps, linear noise schedule
|
|
ae = getae(args).to(device)
|
|
logger.info(f"{model}")
|
|
logger.info(f"DiT Parameters: {sum(p.numel() for p in model.parameters()):,}")
|
|
for n, p in model.named_parameters():
|
|
if p.requires_grad:
|
|
logger.info(f"Training Parameters: {n}")
|
|
|
|
# Setup optimizer (we used default Adam betas=(0.9, 0.999) and a constant learning rate of 1e-4 in our paper):
|
|
opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0)
|
|
|
|
# Setup data:
|
|
dataset = getdataset(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.global_batch_size // dist.get_world_size()),
|
|
shuffle=False,
|
|
sampler=sampler,
|
|
num_workers=args.num_workers,
|
|
pin_memory=True,
|
|
drop_last=True,
|
|
collate_fn=Collate(args)
|
|
)
|
|
logger.info(f"Dataset contains {len(dataset):,} images ({args.data_path})")
|
|
|
|
# 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
|
|
start_time = time()
|
|
|
|
logger.info(f"Training for {args.epochs} epochs...")
|
|
for epoch in range(args.epochs):
|
|
sampler.set_epoch(epoch)
|
|
logger.info(f"Beginning epoch {epoch}...")
|
|
for x, y, attn_mask in loader:
|
|
x = x.to(device)
|
|
y = y.to(device)
|
|
attn_mask = attn_mask.to(device)
|
|
with torch.no_grad():
|
|
x = ae.encode(x) # Map input images to latent space + normalize latents
|
|
t = torch.randint(0, diffusion.num_timesteps, (x.shape[0],), device=device)
|
|
model_kwargs = dict(y=y, attention_mask=attn_mask)
|
|
loss_dict = diffusion.training_losses(model, x, t, model_kwargs)
|
|
loss = loss_dict["loss"].mean()
|
|
opt.zero_grad()
|
|
|
|
loss.backward()
|
|
if args.clip_grad_norm is not None:
|
|
nn.utils.clip_grad_norm_(model.parameters(), args.clip_grad_norm)
|
|
opt.step()
|
|
|
|
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}")
|
|
# Reset monitoring variables:
|
|
running_loss = 0
|
|
log_steps = 0
|
|
start_time = time()
|
|
|
|
# Save DiT 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 DiT-XL/2 with the hyperparameters we used in our paper (except training iters).
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--dataset", type=str, required=True)
|
|
parser.add_argument("--data-path", type=str, required=True)
|
|
parser.add_argument("--results-dir", type=str, default="results")
|
|
parser.add_argument("--model", type=str, choices=list(Diffusion_models.keys()), default="DiT-XL/122")
|
|
parser.add_argument("--num-classes", type=int, default=1000)
|
|
parser.add_argument("--epochs", type=int, default=1400)
|
|
parser.add_argument("--global-batch-size", type=int, default=256)
|
|
parser.add_argument("--global-seed", type=int, default=0)
|
|
parser.add_argument("--num-workers", type=int, default=8)
|
|
parser.add_argument("--log-every", type=int, default=100)
|
|
parser.add_argument("--ckpt-every", type=int, default=50_000)
|
|
|
|
|
|
# --------------------------------------
|
|
parser.add_argument("--ae", type=str, choices=['bair_stride4x2x2', 'ucf101_stride4x4x4',
|
|
'kinetics_stride4x4x4', 'kinetics_stride2x4x4'],
|
|
default="ucf101_stride4x4x4")
|
|
parser.add_argument("--pt-ckpt", type=str, default=None)
|
|
parser.add_argument("--sample-rate", type=int, default=4)
|
|
parser.add_argument("--num-frames", type=int, default=16)
|
|
parser.add_argument("--max-image-size", type=int, default=128)
|
|
parser.add_argument("--dynamic-frames", action="store_true")
|
|
parser.add_argument("--gradient-checkpointing", action="store_true")
|
|
parser.add_argument("--lr", type=float, default=1e-4)
|
|
parser.add_argument("--clip-grad-norm", default=None, type=float, help="the maximum gradient norm (default None)")
|
|
# --------------------------------------
|
|
|
|
|
|
args = parser.parse_args()
|
|
main(args)
|