USER
Here is the code
'''
import argparse
import datetime
import os
import sys
import time
import types
import warnings
from pathlib import Path
current_file_path = Path(__file__).resolve()
sys.path.insert(0, str(current_file_path.parent.parent))
import numpy as np
import torch
from accelerate import Accelerator, InitProcessGroupKwargs
from accelerate.utils import DistributedType
from diffusers.models import AutoencoderKL
from transformers import T5EncoderModel, T5Tokenizer
from mmcv.runner import LogBuffer
from PIL import Image
from torch.utils.data import RandomSampler
from torchvision import transforms
import torch.distributed as dist
from diffusion import IDDPM, DPMS
from diffusion.model.builder import build_model
from diffusion.utils.checkpoint import save_checkpoint, load_checkpoint
from diffusion.utils.dist_utils import synchronize, get_world_size, clip_grad_norm_, flush
from diffusion.utils.logger import get_root_logger, rename_file_with_creation_time
from diffusion.utils.lr_scheduler import build_lr_scheduler
from diffusion.utils.misc import set_random_seed, read_config, init_random_seed, DebugUnderflowOverflow
from diffusion.utils.optimizer import build_optimizer, auto_scale_lr
from diffusion.data.datasets import SimpleDataset
warnings.filterwarnings("ignore") # ignore warning
def set_fsdp_env():
os.environ["ACCELERATE_USE_FSDP"] = 'true'
os.environ["FSDP_AUTO_WRAP_POLICY"] = 'TRANSFORMER_BASED_WRAP'
os.environ["FSDP_BACKWARD_PREFETCH"] = 'BACKWARD_PRE'
def center_crop_arr(pil_image, image_size):
"""
Center cropping implementation from ADM.
https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
"""
while min(*pil_image.size) >= 2 * image_size:
pil_image = pil_image.resize(
tuple(x // 2 for x in pil_image.size), resample=Image.BOX
)
scale = image_size / min(*pil_image.size)
pil_image = pil_image.resize(
tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
)
arr = np.array(pil_image)
crop_y = (arr.shape[0] - image_size) // 2
crop_x = (arr.shape[1] - image_size) // 2
return Image.fromarray(arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size])
def train():
if config.get('debug_nan', False):
DebugUnderflowOverflow(model)
logger.info('NaN debugger registered. Start to detect overflow during training.')
time_start, last_tic = time.time(), time.time()
log_buffer = LogBuffer()
global_step = start_step + 1
load_vae_feat = False #getattr(train_dataloader.dataset, 'load_vae_feat', False)
load_t5_feat = False #getattr(train_dataloader.dataset, 'load_t5_feat', False)
# Now you train the model
for epoch in range(start_epoch + 1, config.num_epochs + 1):
data_time_start= time.time()
data_time_all = 0
loss_sum = 0.
for step, batch in enumerate(train_dataloader):
if step < skip_step:
global_step += 1
continue # skip data in the resumed ckpt
if load_vae_feat:
z = batch[0]
else:
with torch.no_grad():
with torch.cuda.amp.autocast(enabled=(config.mixed_precision == 'fp16' or config.mixed_precision == 'bf16')):
posterior = vae.encode(batch[0]).latent_dist
if config.sample_posterior:
z = posterior.sample()
else:
z = posterior.mode()
clean_images = z * config.scale_factor
data_info = None # batch[3]
if load_t5_feat:
y = batch[1]
y_mask = batch[2]
else:
with torch.no_grad():
txt_tokens = tokenizer(
batch[1], max_length=max_length, padding="max_length", truncation=True, return_tensors="pt"
).to(accelerator.device)
y = text_encoder(
txt_tokens.input_ids, attention_mask=txt_tokens.attention_mask)[0][:, None]
y_mask = txt_tokens.attention_mask[:, None, None]
# Sample a random timestep for each image
bs = clean_images.shape[0]
timesteps = torch.randint(0, config.train_sampling_steps, (bs,), device=clean_images.device).long()
grad_norm = None
data_time_all += time.time() - data_time_start
with accelerator.accumulate(model):
# Predict the noise residual
optimizer.zero_grad()
loss_term = train_diffusion.training_losses(model, clean_images, timesteps, model_kwargs=dict(y=y, mask=y_mask, data_info=data_info))
loss = loss_term['loss']
loss = torch.where(torch.isnan(loss), torch.zeros_like(loss), loss)
loss = loss.mean()
# loss = torch.nan_to_num(loss)
# if not torch.isnan(loss):
accelerator.backward(loss)
loss_sum += loss.item()
if accelerator.sync_gradients:
grad_norm = accelerator.clip_grad_norm_(model.parameters(), config.gradient_clip)
optimizer.step()
lr_scheduler.step()
lr = lr_scheduler.get_last_lr()[0]
logs = {args.loss_report_name: accelerator.gather(loss).mean().item()}
logs.update(avg_loss=loss_sum / (step + 1))
if grad_norm is not None:
logs.update(grad_norm=accelerator.gather(grad_norm).mean().item())
log_buffer.update(logs)
if (step + 1) % config.log_interval == 0 or (step + 1) == 1:
t = (time.time() - last_tic) / config.log_interval
t_d = data_time_all / config.log_interval
avg_time = (time.time() - time_start) / (global_step + 1)
eta = str(datetime.timedelta(seconds=int(avg_time * (total_steps - global_step - 1))))
eta_epoch = str(datetime.timedelta(seconds=int(avg_time * (len(train_dataloader) - step - 1))))
log_buffer.average()
info = f"Step/Epoch [{global_step}/{epoch}][{step + 1}/{len(train_dataloader)}]:total_eta: {eta}, " \
f"epoch_eta:{eta_epoch}, time_all:{t:.3f}, time_data:{t_d:.3f}, lr:{lr:.3e}, s:({model.module.h}, {model.module.w}), "
info += ', '.join([f"{k}:{v:.4f}" for k, v in log_buffer.output.items()])
logger.info(info)
last_tic = time.time()
log_buffer.clear()
data_time_all = 0
logs.update(lr=lr)
accelerator.log(logs, step=global_step)
global_step += 1
data_time_start = time.time()
if global_step % config.save_model_steps == 0:
accelerator.wait_for_everyone()
if accelerator.is_main_process:
os.umask(0o000)
save_checkpoint(os.path.join(config.work_dir, 'checkpoints'),
epoch=epoch,
step=global_step,
model=accelerator.unwrap_model(model),
optimizer=optimizer,
lr_scheduler=lr_scheduler
)
if config.visualize and (global_step % config.eval_sampling_steps == 0 or (step + 1) == 1):
accelerator.wait_for_everyone()
if epoch % config.save_model_epochs == 0 or epoch == config.num_epochs:
accelerator.wait_for_everyone()
if accelerator.is_main_process:
os.umask(0o000)
save_checkpoint(os.path.join(config.work_dir, 'checkpoints'),
epoch=epoch,
step=global_step,
model=accelerator.unwrap_model(model),
optimizer=optimizer,
lr_scheduler=lr_scheduler
)
accelerator.wait_for_everyone()
def parse_args():
parser = argparse.ArgumentParser(description="Process some integers.")
parser.add_argument("config", type=str, help="config")
parser.add_argument("--cloud", action='store_true', default=False, help="cloud or local machine")
parser.add_argument('--work-dir', help='the dir to save logs and models')
parser.add_argument('--resume-from', help='the dir to resume the training')
parser.add_argument('--load-from', default=None, help='the dir to load a ckpt for training')
parser.add_argument('--local-rank', type=int, default=-1)
parser.add_argument('--local_rank', type=int, default=-1)
parser.add_argument('--debug', action='store_true')
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(
"--tracker_project_name",
type=str,
default="text2image-fine-tune",
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("--loss_report_name", type=str, default="loss")
args = parser.parse_args()
return args
if __name__ == '__main__':
args = parse_args()
config = read_config(args.config)
if args.work_dir is not None:
config.work_dir = args.work_dir
if args.debug:
config.log_interval = 1
config.train_batch_size = 2
os.umask(0o000)
os.makedirs(config.work_dir, exist_ok=True)
init_handler = InitProcessGroupKwargs()
init_handler.timeout = datetime.timedelta(seconds=5400) # change timeout to avoid a strange NCCL bug
# Initialize accelerator and tensorboard logging
if config.use_fsdp:
init_train = 'FSDP'
from accelerate import FullyShardedDataParallelPlugin
from torch.distributed.fsdp.fully_sharded_data_parallel import FullStateDictConfig
set_fsdp_env()
fsdp_plugin = FullyShardedDataParallelPlugin(state_dict_config=FullStateDictConfig(offload_to_cpu=False, rank0_only=False),)
else:
init_train = 'DDP'
fsdp_plugin = None
even_batches = True
if config.multi_scale:
even_batches=False,
accelerator = Accelerator(
mixed_precision=config.mixed_precision,
gradient_accumulation_steps=config.gradient_accumulation_steps,
log_with=args.report_to,
project_dir=os.path.join(config.work_dir, "logs"),
fsdp_plugin=fsdp_plugin,
even_batches=even_batches,
kwargs_handlers=[init_handler]
)
log_name = 'train_log.log'
if accelerator.is_main_process:
if os.path.exists(os.path.join(config.work_dir, log_name)):
rename_file_with_creation_time(os.path.join(config.work_dir, log_name))
logger = get_root_logger(os.path.join(config.work_dir, log_name))
logger.info(accelerator.state)
config.seed = 2024 # init_random_seed(config.get('seed', None))
set_random_seed(config.seed)
if accelerator.is_main_process:
config.dump(os.path.join(config.work_dir, 'config.py'))
logger.info(f"Config: \n{config.pretty_text}")
logger.info(f"World_size: {get_world_size()}, seed: {config.seed}")
logger.info(f"Initializing: {init_train} for training")
image_size = config.image_size # @param [256, 512]
latent_size = int(image_size) // 8
pred_sigma = getattr(config, 'pred_sigma', True)
learn_sigma = getattr(config, 'learn_sigma', True) and pred_sigma
max_length = config.model_max_length
kv_compress_config = config.kv_compress_config if config.kv_compress else None
vae = None
# if not config.data.load_vae_feat:
vae = AutoencoderKL.from_pretrained(config.vae_pretrained, torch_dtype=torch.float32).to(accelerator.device)
config.scale_factor = vae.config.scaling_factor
tokenizer = text_encoder = None
pipeline_load_from = config.t5_path
tokenizer = T5Tokenizer.from_pretrained(pipeline_load_from)
text_encoder = T5EncoderModel.from_pretrained(pipeline_load_from, torch_dtype=torch.float16).to(accelerator.device)
logger.info(f"vae scale factor: {config.scale_factor}")
config.visualize = False
if config.visualize:
# preparing embeddings for visualization. We put it here for saving GPU memory
validation_prompts = [
"dog",
"portrait photo of a girl, photograph, highly detailed face, depth of field",
"Self-portrait oil painting, a beautiful cyborg with golden hair, 8k",
"Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
"A photo of beautiful mountain with realistic sunset and blue lake, highly detailed, masterpiece",
]
null_tokens = tokenizer(
"", max_length=max_length, padding="max_length", truncation=True, return_tensors="pt"
).to(accelerator.device)
null_token_emb = text_encoder(null_tokens.input_ids, attention_mask=null_tokens.attention_mask)[0]
torch.save(
{'uncond_prompt_embeds': null_token_emb, 'uncond_prompt_embeds_mask': null_tokens.attention_mask},
f'ckpts/null_embed_diffusers_{max_length}token.pth')
flush()
model_kwargs = {"pe_interpolation": config.pe_interpolation, "config": config,
"model_max_length": max_length, "qk_norm": config.qk_norm,
"kv_compress_config": kv_compress_config, "micro_condition": config.micro_condition}
# build models
train_diffusion = IDDPM(str(config.train_sampling_steps), learn_sigma=learn_sigma, pred_sigma=pred_sigma, snr=config.snr_loss)
model = build_model(config.model,
config.grad_checkpointing,
config.get('fp32_attention', False),
input_size=latent_size,
learn_sigma=learn_sigma,
pred_sigma=pred_sigma,
**model_kwargs).train()
logger.info(f"{model.__class__.__name__} Model Parameters: {sum(p.numel() for p in model.parameters()):,}")
if args.load_from is not None:
config.load_from = args.load_from
if config.load_from is not None:
missing, unexpected = load_checkpoint(
config.load_from, model, load_ema=config.get('load_ema', False), max_length=max_length)
logger.warning(f'Missing keys: {missing}')
logger.warning(f'Unexpected keys: {unexpected}')
# prepare for FSDP clip grad norm calculation
if accelerator.distributed_type == DistributedType.FSDP:
for m in accelerator._models:
m.clip_grad_norm_ = types.MethodType(clip_grad_norm_, m)
"""
for simple image-text dataset with json
"""
transform = transforms.Compose([
transforms.Resize(image_size),
transforms.Lambda(lambda pil_image: center_crop_arr(pil_image, image_size)),
# transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
])
world_size = get_world_size()
data_path = config.data_path
dataset = SimpleDataset(path=data_path, transform=transform)
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
sampler = DistributedSampler(
dataset,
num_replicas=dist.get_world_size(),
rank=dist.get_rank(),
shuffle=True,
seed=config.seed,
)
train_dataloader = DataLoader(
dataset,
batch_size=int(config.train_batch_size),
shuffle=True,
# sampler=sampler,
num_workers=config.num_workers,
pin_memory=True,
drop_last=True
)
print(len(train_dataloader))
# build optimizer and lr scheduler
lr_scale_ratio = 1
if config.get('auto_lr', None):
lr_scale_ratio = auto_scale_lr(config.train_batch_size * get_world_size() * config.gradient_accumulation_steps,
config.optimizer, **config.auto_lr)
optimizer = build_optimizer(model, config.optimizer)
lr_scheduler = build_lr_scheduler(config, optimizer, train_dataloader, lr_scale_ratio)
timestamp = time.strftime("%Y-%m-%d_%H:%M:%S", time.localtime())
if accelerator.is_main_process:
tracker_config = dict(vars(config))
try:
accelerator.init_trackers(args.tracker_project_name, tracker_config)
except:
accelerator.init_trackers(f"tb_{timestamp}")
start_epoch = 0
start_step = 0
skip_step = config.skip_step
total_steps = len(train_dataloader) * config.num_epochs
model = accelerator.prepare(model)
optimizer, train_dataloader, lr_scheduler = accelerator.prepare(optimizer, train_dataloader, lr_scheduler)
train()
'''