Spaces:
Running
on
Zero
Running
on
Zero
File size: 6,962 Bytes
b6a9b6d |
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 |
# 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.
"""
import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch
# the first flag below was False when we tested this script but True makes A100 training a lot faster:
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
from torch.utils.data import DataLoader
import numpy as np
from copy import deepcopy
from glob import glob
from time import time
import argparse
import os
import yaml
from accelerate import Accelerator
from torch.utils.tensorboard import SummaryWriter
from core.models import DiT_models
from core.diffusion import create_diffusion
from core.dataset import ImageParamsDataset
from core.utils.train_utils import create_logger, update_ema, requires_grad
#################################################################################
# Training Loop #
#################################################################################
def main(cfg):
"""
Trains a new DiT model.
"""
assert torch.cuda.is_available(), "Training currently requires at least one GPU."
# Setup accelerator:
accelerator = Accelerator()
device = accelerator.device
# Setup an experiment folder:
if accelerator.is_main_process:
os.makedirs(cfg["save_dir"], exist_ok=True) # Make results folder (holds all experiment subfolders)
save_dir = cfg["save_dir"]
experiment_index = len(glob(f"{save_dir}/*"))
experiment_dir = "{}/{:03d}-{}-{}-{}".format(save_dir, experiment_index, cfg["model"], cfg["epochs"], cfg["batch_size"]) # Create an experiment folder
checkpoint_dir = "{}/checkpoints".format(experiment_dir) # 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}")
writer = SummaryWriter(experiment_dir)
# Create model:
latent_size = cfg["num_params"]
condition_channels = 768
model = DiT_models[cfg["model"]](input_size=latent_size, condition_channels=condition_channels)
# Note that parameter initialization is done within the DiT constructor
model = model.to(device)
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
if accelerator.is_main_process:
logger.info(f"DiT Parameters: {sum(p.numel() for p in model.parameters()):,}")
# Setup optimizer (we used default Adam betas=(0.9, 0.999) and a constant learning rate of 1e-4 in our paper):
optimizer = torch.optim.AdamW(model.parameters(), lr=float(cfg["lr"]), weight_decay=0)
# Setup data:
dataset = ImageParamsDataset(cfg["data_root"], cfg["train_file"], cfg["params_dict_file"])
loader = DataLoader(
dataset,
batch_size=int(cfg["batch_size"] // accelerator.num_processes),
shuffle=True,
num_workers=cfg["num_workers"],
pin_memory=True,
drop_last=True
)
if accelerator.is_main_process:
logger.info(f"Dataset contains {len(dataset):,} images")
# Prepare models for training:
update_ema(ema, model, 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
model, optimizer, loader = accelerator.prepare(model, optimizer, loader)
# Variables for monitoring/logging purposes:
train_steps = 0
log_steps = 0
running_loss = 0
start_time = time()
if accelerator.is_main_process:
logger.info("Training for {} epochs...".format(cfg["epochs"]))
# main training loop
for epoch in range(int(cfg["epochs"])):
if accelerator.is_main_process:
logger.info(f"Beginning epoch {epoch}...")
for x, img_feat, img in loader:
# prepare the inputs
x = x.to(device)
img_feat = img_feat.to(device)
x = x.unsqueeze(dim=1) # [B, 1, N]
t = torch.randint(0, diffusion.num_timesteps, (x.shape[0],), device=device)
model_kwargs = dict(y=img_feat)
loss_dict = diffusion.training_losses(model, x, t, model_kwargs)
loss = loss_dict["loss"].mean()
optimizer.zero_grad()
accelerator.backward(loss)
optimizer.step()
update_ema(ema, model)
writer.add_scalar("train/loss", loss.item(), train_steps)
# Log loss values:
running_loss += loss.item()
log_steps += 1
train_steps += 1
if train_steps % cfg["logging_iter"] == 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)
avg_loss = avg_loss.item() / accelerator.num_processes
if accelerator.is_main_process:
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 % cfg["ckpt_iter"] == 0 and train_steps > 0:
if accelerator.is_main_process:
checkpoint = {
"model": model.state_dict(),
"ema": ema.state_dict(),
"optimizer": optimizer.state_dict(),
"config": cfg,
}
checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}.pt"
torch.save(checkpoint, checkpoint_path)
logger.info(f"Saved checkpoint to {checkpoint_path}")
model.eval() # important! This disables randomized embedding dropout
# do any sampling/FID calculation/etc. with ema (or model) in eval mode ...
if accelerator.is_main_process:
writer.flush()
logger.info("Done!")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--config", type=str, required=True)
args = parser.parse_args()
with open(args.config) as f:
cfg = yaml.load(f, Loader=yaml.FullLoader)
main(cfg)
|