wehub-resource-sync
221 lines
9.7 KiB
Python
221 lines
9.7 KiB
Python
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
import math
|
|
import numpy as np
|
|
from tqdm import tqdm
|
|
import torch
|
|
import torch.distributed as dist
|
|
from accelerate.utils import set_seed
|
|
|
|
from safetensors.torch import load_file
|
|
from tokenizer_models import AutoencoderKL, load_vae
|
|
|
|
from schedule.dpm_solver import DPMSolverMultistepScheduler
|
|
from models import All_models
|
|
from utils import safe_blob_dump
|
|
from metrics import compute_fid_without_store, compute_inception_score_from_tensor
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description="Simple example of a training script.")
|
|
parser.add_argument(
|
|
"--seed",
|
|
type=int,
|
|
default=0,
|
|
help="A seed to use for the random number generator. Can be negative to not set a seed.",
|
|
)
|
|
parser.add_argument(
|
|
"--model",
|
|
type=str,
|
|
default="Transformer-L",
|
|
help="The config of the UNet model to train, leave as None to use standard DDPM configuration.",
|
|
)
|
|
parser.add_argument(
|
|
"--vae",
|
|
type=str,
|
|
default=None,
|
|
)
|
|
parser.add_argument(
|
|
"--train_data_dir",
|
|
type=str,
|
|
default="/tmp/ILSVRC/Data/CLS-LOC/train",
|
|
help=(
|
|
"A folder containing the training data. Folder contents must follow the structure described in"
|
|
" https://huggingface.co/docs/datasets/image_dataset#imagefolder. In particular, a `metadata.jsonl` file"
|
|
" must exist to provide the captions for the images. Ignored if `dataset_name` is specified."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--ref_stat_path",
|
|
type=str,
|
|
default="/mnt/unilm/hangbo/beit3/t2i/assets/fid_stats/imagenet_256_val.npz",
|
|
)
|
|
parser.add_argument(
|
|
"--image_size",
|
|
type=int,
|
|
default=256,
|
|
help=(
|
|
"The image_size for input images, all the images in the train/validation dataset will be resized to this"
|
|
" image_size"
|
|
),
|
|
)
|
|
parser.add_argument("--num-classes", type=int, default=1000)
|
|
parser.add_argument(
|
|
"--mixed_precision",
|
|
type=str,
|
|
default="no",
|
|
choices=["no", "fp16", "bf16"],
|
|
help=(
|
|
"Whether to use mixed precision. Choose"
|
|
"between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
|
|
"and an Nvidia Ampere GPU."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--batch_size", type=int, default=32, help="Batch size (per device) for the training dataloader."
|
|
)
|
|
parser.add_argument(
|
|
"--steps_per_class", type=int, default=50, help="Number of steps per class."
|
|
)
|
|
parser.add_argument("--force_diffusion", action="store_true", help="Whether to force the use of diffusion models.")
|
|
parser.add_argument("--use_ema", action="store_true", help="Whether to use Exponential Moving Average for the final model weights.")
|
|
parser.add_argument("--ddpm_num_steps", type=int, default=1000)
|
|
parser.add_argument("--ddpm_num_inference_steps", type=int, default=250)
|
|
parser.add_argument("--ddpm_beta_schedule", type=str, default="cosine", help="The beta schedule to use for DDPM.")
|
|
parser.add_argument("--prediction_type", type=str, default="epsilon", help="Whether the model should predict the 'epsilon'/noise error or directly the reconstructed image 'x0'.")
|
|
parser.add_argument("--cfg-scale", type=float, default=4.0)
|
|
parser.add_argument(
|
|
"--checkpoint",
|
|
type=str,
|
|
default=None,
|
|
help=(
|
|
"Whether training should be resumed from a previous checkpoint. Use a path saved by"
|
|
' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
|
|
),
|
|
)
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
def suppress_output(rank):
|
|
"""Suppress output for all processes except the one with rank 0."""
|
|
if rank != 0:
|
|
sys.stdout = open(os.devnull, 'w')
|
|
|
|
@torch.no_grad()
|
|
def main(args):
|
|
set_seed(args.seed)
|
|
dist.init_process_group(backend="gloo", init_method='env://')
|
|
rank = dist.get_rank()
|
|
suppress_output(rank)
|
|
print(args)
|
|
device = f"cuda:{rank}" if torch.cuda.is_available() else "cpu"
|
|
if args.mixed_precision == "bf16":
|
|
dtype = torch.bfloat16
|
|
elif args.mixed_precision == "fp16":
|
|
dtype = torch.float16
|
|
else:
|
|
dtype = torch.float32
|
|
prefix = "ema" if args.use_ema else "standard"
|
|
exp_name = f"{prefix}_{args.steps_per_class}_{args.cfg_scale}_{args.ddpm_beta_schedule}_{args.ddpm_num_inference_steps}"
|
|
print(f"Exp_name {exp_name}")
|
|
vae, input_size, latent_size, flatten_input = load_vae(args.vae, args.image_size)
|
|
|
|
vae.eval()
|
|
other_state = torch.load(os.path.join(args.checkpoint, "other_state.pth"), map_location="cpu")
|
|
scaling_factor = other_state["scaling_factor"]
|
|
bias_factor = other_state["bias_factor"]
|
|
print(f"Scaling factor: {scaling_factor}, Bias factor: {bias_factor}")
|
|
# Potentially load in the weights and states from a previous save
|
|
latent_path = os.path.join(args.checkpoint, f"latent_{exp_name}.pth")
|
|
if os.path.exists(latent_path) and not args.force_diffusion:
|
|
all_latent_gather = torch.load(latent_path)
|
|
print("Loaded latent from file.")
|
|
else:
|
|
model = All_models[args.model](
|
|
input_size=input_size,
|
|
in_channels=latent_size,
|
|
num_classes=args.num_classes,
|
|
flatten_input=flatten_input,
|
|
).to(device).to(dtype)
|
|
noise_scheduler = DPMSolverMultistepScheduler(num_train_timesteps=args.ddpm_num_steps, beta_schedule=args.ddpm_beta_schedule, prediction_type=args.prediction_type)
|
|
model.eval()
|
|
if args.checkpoint:
|
|
if args.use_ema and other_state["ema"] is not None:
|
|
checkpoint = other_state["ema"]["shadow_params"]
|
|
for model_param, ema_param in zip(model.parameters(), checkpoint):
|
|
model_param.data = ema_param.data.to(device).to(dtype)
|
|
print(f"Loaded model from checkpoint {args.checkpoint}, EMA applied.")
|
|
else:
|
|
if os.path.exists(os.path.join(args.checkpoint, "model.safetensors")):
|
|
checkpoint = load_file(os.path.join(args.checkpoint, "model.safetensors"))
|
|
elif os.path.exists(os.path.join(args.checkpoint, "pytorch_model")):
|
|
checkpoint = torch.load(os.path.join(args.checkpoint, "pytorch_model", "mp_rank_00_model_states.pt"), map_location="cpu")["module"]
|
|
|
|
model.load_state_dict(checkpoint)
|
|
print(f"Loaded model from checkpoint {args.checkpoint}.")
|
|
|
|
def p_sample(model, image):
|
|
noise_scheduler.set_timesteps(args.ddpm_num_inference_steps)
|
|
for t in noise_scheduler.timesteps:
|
|
model_output = model(image, t.repeat(image.shape[0]).to(image))
|
|
image = noise_scheduler.step(model_output, t, image).prev_sample
|
|
return image
|
|
|
|
all_latent = []
|
|
class_start, class_end = args.num_classes // dist.get_world_size() * rank, args.num_classes // dist.get_world_size() * (rank + 1)
|
|
classes = torch.arange(class_start, class_end, device=device).repeat(args.steps_per_class)
|
|
classes = classes.chunk(math.ceil(classes.size(0) / args.batch_size))
|
|
for y in tqdm(classes, disable=rank != 0):
|
|
y_null = torch.full_like(y, args.num_classes, device=device)
|
|
y = torch.cat([y, y_null], 0)
|
|
# Sample images:
|
|
samples = model.sample_with_cfg(y, args.cfg_scale, p_sample)
|
|
all_latent.append(samples.float().cpu())
|
|
|
|
all_latent = torch.cat(all_latent, 0)
|
|
all_latent_gather = [torch.zeros_like(all_latent) for _ in range(dist.get_world_size())]
|
|
dist.all_gather(all_latent_gather, all_latent)
|
|
all_latent_gather = torch.cat(all_latent_gather, 0)
|
|
if rank == 0:
|
|
torch.save(all_latent_gather, latent_path)
|
|
|
|
if rank == 0:
|
|
all_images = torch.zeros((all_latent_gather.size(0), 3, 256, 256))
|
|
if args.image_size != 256:
|
|
transform = torch.nn.Upsample(size=(256, 256), mode="bilinear")
|
|
else:
|
|
transform = torch.nn.Identity()
|
|
idx = 0
|
|
for samples in tqdm(all_latent_gather.chunk(math.ceil(all_latent_gather.size(0) / args.batch_size))):
|
|
images = vae.decode(samples.to(device).to(dtype) / scaling_factor - bias_factor)
|
|
images = transform(images)
|
|
images = (torch.clamp(images.float(), -1.0, 1.0) * 0.5 + 0.5).cpu().float()
|
|
all_images[idx:idx + images.shape[0]] = images
|
|
idx += images.shape[0]
|
|
|
|
print(all_images.shape)
|
|
fid_score = compute_fid_without_store(all_images, args.ref_stat_path, batch_size=args.batch_size, device=device)
|
|
print(fid_score)
|
|
IS_mean, IS_std = compute_inception_score_from_tensor(
|
|
all_images,
|
|
batch_size=args.batch_size,
|
|
device=device,
|
|
)
|
|
print(IS_mean, IS_std)
|
|
result_path = os.path.join(args.checkpoint, f"result_{exp_name}.json")
|
|
result = {
|
|
"fid": fid_score.item(),
|
|
"IS_mean": IS_mean.item(),
|
|
"IS_std": IS_std.item(),
|
|
}
|
|
safe_blob_dump(result_path, result)
|
|
image_path = os.path.join(args.checkpoint, f"images_{exp_name}.npz")
|
|
all_images = (all_images * 255.0).clamp(0, 255).to(torch.uint8).permute(0, 2, 3, 1).numpy()
|
|
np.savez_compressed(image_path, all_images)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = parse_args()
|
|
main(args) |