nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[refactor] adapt to old training code

Browse Source 
Former-commit-id: 358d362ecbc100b21ea02895cb0426cf5b386a72
This commit is contained in:
qqingzheng
2024-03-06 13:58:13 +00:00
parent 1bed11aee3
commit 2b272863b0
2 changed files with 296 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
opensora/models/ae/__init__.py
+2 -2
View File
@@ -1,5 +1,5 @@
from .imagebase import imagebase_ae
from .videobase import videobase_ae
from .imagebase import imagebase_ae, vae, vqvae
from .videobase import videobase_ae, videovae, videovqvae
from .videobase import (
VideoGPTConfiguration,
VideoGPTVQVAE
opensora/train/train.py
+294
View File
@@ -0,0 +1,294 @@
# 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)