def save_model(mp_trainer, opt, step):
if dist.get_rank() == 0:
th.save(
mp_trainer.master_params_to_state_dict(mp_trainer.master_params),
os.path.join(logger.get_dir(), f"model{step:06d}.pt"),
)
th.save(opt.state_dict(),
os.path.join(logger.get_dir(), f"opt{step:06d}.pt"))
def main():
args = create_argparser().parse_args()
pprint({k:v for k,v in args.__dict__.items()})
dist_util.setup_dist()
logger.configure()
logger.log("creating model...")
pprint(args_to_dict(args, sr_model_and_diffusion_defaults().keys()))
model, diffusion = sr_create_model_and_diffusion(
**args_to_dict(args, sr_model_and_diffusion_defaults().keys())
)
# skips
# load_tolerant(model, args.model_path)
model.load_state_dict(
dist_util.load_state_dict(args.model_path, map_location="cpu")
)
model.to(dist_util.dev())
if args.use_fp16:
model.convert_to_fp16()
model.eval()
logger.log("loading data...")
data = load_data_for_worker(args.base_samples, args.batch_size, args.class_cond)
logger.log("creating samples...")
all_images = []
while len(all_images) * args.batch_size < args.num_samples:
model_kwargs = next(data)
model_kwargs = {k: v.to(dist_util.dev()) for k, v in model_kwargs.items()}
sample = diffusion.p_sample_loop(
model,
(args.batch_size, 3, args.large_size, args.large_size),
clip_denoised=args.clip_denoised,
model_kwargs=model_kwargs,
)
sample = ((sample + 1) * 127.5).clamp(0, 255).to(th.uint8)
sample = sample.permute(0, 2, 3, 1)
sample = sample.contiguous()
all_samples = [th.zeros_like(sample) for _ in range(dist.get_world_size())]
dist.all_gather(all_samples, sample) # gather not supported with NCCL
for sample in all_samples:
all_images.append(sample.cpu().numpy())
logger.log(f"created {len(all_images) * args.batch_size} samples")
arr = np.concatenate(all_images, axis=0)
arr = arr[: args.num_samples]
if dist.get_rank() == 0:
shape_str = "x".join([str(x) for x in arr.shape])
out_path = os.path.join(logger.get_dir(), f"samples_{shape_str}.npz")
logger.log(f"saving to {out_path}")
np.savez(out_path, arr)
dist.barrier()
logger.log("sampling complete")
def run_bpd_evaluation(model, diffusion, data, num_samples, clip_denoised):
all_bpd = []
all_metrics = {"vb": [], "mse": [], "xstart_mse": []}
num_complete = 0
while num_complete < num_samples:
batch, model_kwargs = next(data)
batch = batch.to(dist_util.dev())
model_kwargs = {
k: v.to(dist_util.dev())
for k, v in model_kwargs.items()
}
minibatch_metrics = diffusion.calc_bpd_loop(
model,
batch,
clip_denoised=clip_denoised,
model_kwargs=model_kwargs)
for key, term_list in all_metrics.items():
terms = minibatch_metrics[key].mean(dim=0) / dist.get_world_size()
dist.all_reduce(terms)
term_list.append(terms.detach().cpu().numpy())
total_bpd = minibatch_metrics["total_bpd"]
total_bpd = total_bpd.mean() / dist.get_world_size()
dist.all_reduce(total_bpd)
all_bpd.append(total_bpd.item())
num_complete += dist.get_world_size() * batch.shape[0]
logger.log(f"done {num_complete} samples: bpd={np.mean(all_bpd)}")
if dist.get_rank() == 0:
for name, terms in all_metrics.items():
out_path = os.path.join(logger.get_dir(), f"{name}_terms.npz")
logger.log(f"saving {name} terms to {out_path}")
np.savez(out_path, np.mean(np.stack(terms), axis=0))
dist.barrier()
logger.log("evaluation complete")
def main():
args = create_argparser().parse_args()
dist_util.setup_dist()
logger.configure()
logger.log("creating model and diffusion...")
model, diffusion = create_model_and_diffusion(
**args_to_dict(args, model_and_diffusion_defaults().keys())
)
model.load_state_dict(
dist_util.load_state_dict(args.model_path, map_location="cpu")
)
model.to(dist_util.dev())
if args.use_fp16:
model.convert_to_fp16()
model.eval()
logger.log("sampling...")
all_images = []
all_labels = []
while len(all_images) * args.batch_size < args.num_samples:
model_kwargs = {}
if args.class_cond:
classes = th.randint(
low=0, high=NUM_CLASSES, size=(args.batch_size,), device=dist_util.dev()
)
model_kwargs["y"] = classes
sample_fn = (
diffusion.p_sample_loop if not args.use_ddim else diffusion.ddim_sample_loop
)
sample = sample_fn(
model,
(args.batch_size, 3, args.image_size, args.image_size),
clip_denoised=args.clip_denoised,
model_kwargs=model_kwargs,
)
sample = ((sample + 1) * 127.5).clamp(0, 255).to(th.uint8)
sample = sample.permute(0, 2, 3, 1)
sample = sample.contiguous()
gathered_samples = [th.zeros_like(sample) for _ in range(dist.get_world_size())]
dist.all_gather(gathered_samples, sample) # gather not supported with NCCL
all_images.extend([sample.cpu().numpy() for sample in gathered_samples])
if args.class_cond:
gathered_labels = [
th.zeros_like(classes) for _ in range(dist.get_world_size())
]
dist.all_gather(gathered_labels, classes)
all_labels.extend([labels.cpu().numpy() for labels in gathered_labels])
logger.log(f"created {len(all_images) * args.batch_size} samples")
arr = np.concatenate(all_images, axis=0)
arr = arr[: args.num_samples]
if args.class_cond:
label_arr = np.concatenate(all_labels, axis=0)
label_arr = label_arr[: args.num_samples]
if dist.get_rank() == 0:
shape_str = "x".join([str(x) for x in arr.shape])
out_path = os.path.join(logger.get_dir(), f"samples_{shape_str}.npz")
logger.log(f"saving to {out_path}")
if args.class_cond:
np.savez(out_path, arr, label_arr)
else:
np.savez(out_path, arr)
dist.barrier()
logger.log("sampling complete")
def main(**kwargs):
args = create_argparser().parse_args()
kw = {k: v for k, v in kwargs.items() if k in args}
args.__dict__.update(**kw)
dist_util.setup_dist()
logger.configure()
logger.log("creating model and diffusion...")
pprint(args_to_dict(args, model_and_diffusion_defaults().keys()))
model, diffusion = create_model_and_diffusion(
**args_to_dict(args,
model_and_diffusion_defaults().keys()))
print(f"loading state dict {args.model_path} -> model")
model.load_state_dict(
dist_util.load_state_dict(args.model_path, map_location="cpu"))
model.to(dist_util.dev())
if args.use_fp16:
model.convert_to_fp16()
model.eval()
logger.log("loading classifier...")
classifier = create_classifier(
**args_to_dict(args,
classifier_defaults().keys()))
print(f"loading state dict {args.classifier_path}")
classifier.load_state_dict(
dist_util.load_state_dict(args.classifier_path, map_location="cpu"))
classifier.to(dist_util.dev())
if args.classifier_use_fp16:
classifier.convert_to_fp16()
classifier.eval()
def cond_fn(x, t, y=None):
assert y is not None
with th.enable_grad():
# print(f" cond_fn(x,t,y) t {t}")
# t, batch_size time sampling 999 -> 0
x_in = x.detach().requires_grad_(True)
logits = classifier(x_in, t)
# print(f" .. x:{x_in.shape}, t:{t.shape}, logits:{logits.shape}, y: {y.shape}")
# x:(batch_size, channels, image_size, image_size), t:(batch_size), (batch_size, numclasses), y: (batch_size)
log_probs = F.log_softmax(logits, dim=-1)
selected = log_probs[range(len(logits)), y.view(-1)]
# print(f" .. selected, softmax(logits)[range(), y] {selected}") # (batch_size) floats
cond = th.autograd.grad(selected.sum(),
x_in)[0] * args.classifier_scale
# print(f" .. cond: {tuple(cond.shape)}, args.classifier_scale {args.classifier_scale}")
# cond: (batch_size, 3, image_size, image_size), args.classifier_scale 0.5
return cond
def logt(x):
if isinstance(x, th.Tensor):
out = f" {tuple(x.shape)}"
if x.ndim == 1:
out += f"{x}"
elif isinstance(x, (int, float)):
out = f" {x}"
return out
def model_fn(x, t, y=None):
assert y is not None
print(f"timestep {t.tolist()} conditional y {y.tolist()}")
#print(f"model_fn, x {logt(x)}, t {logt(t)} y {logt(y)}")
return model(x, t, y if args.class_cond else None)
logger.log("sampling...")
all_images = []
all_labels = []
while len(all_images) * args.batch_size < args.num_samples:
model_kwargs = {}
classes = th.randint(low=0,
high=NUM_CLASSES,
size=(args.batch_size, ),
device=dist_util.dev())
model_kwargs["y"] = classes
if args.use_ddim:
print("sample_fn = diffusion.ddim_sample_loop: args.use_ddim")
else:
print("sample_fn = diffusion.p_sample_loop: not args.use_ddim")
sample_fn = (diffusion.p_sample_loop
if not args.use_ddim else diffusion.ddim_sample_loop)
# print(f"sample_fn args.batch_size {args.batch_size}, args.image_size {args.image_size} args.clip_denoised {args.clip_denoised} model_kwargs {model_kwargs}")
# model_kwargs['y']: class conditioner e.g [ 53, 37, 609, 498, 679, 38, 242, 705, 253, 822, 721, 762, 64, 42, 337, 483]
sample = sample_fn(
model_fn,
(args.batch_size, 3, args.image_size, args.image_size),
clip_denoised=args.clip_denoised,
model_kwargs=model_kwargs,
cond_fn=cond_fn,
device=dist_util.dev(),
)
sample = ((sample + 1) * 127.5).clamp(0, 255).to(th.uint8)
sample = sample.permute(0, 2, 3, 1)
sample = sample.contiguous()
gathered_samples = [
th.zeros_like(sample) for _ in range(dist.get_world_size())
]
dist.all_gather(gathered_samples,
sample) # gather not supported with NCCL
all_images.extend(
[sample.cpu().numpy() for sample in gathered_samples])
gathered_labels = [
th.zeros_like(classes) for _ in range(dist.get_world_size())
]
dist.all_gather(gathered_labels, classes)
all_labels.extend([labels.cpu().numpy() for labels in gathered_labels])
logger.log(f"created {len(all_images) * args.batch_size} samples")
arr = np.concatenate(all_images, axis=0)
arr = arr[:args.num_samples]
label_arr = np.concatenate(all_labels, axis=0)
label_arr = label_arr[:args.num_samples]
if dist.get_rank() == 0:
shape_str = "x".join([str(x) for x in arr.shape])
out_path = os.path.join(logger.get_dir(), f"samples_{shape_str}.npz")
logger.log(f"saving to {out_path}")
np.savez(out_path, arr, label_arr)
dist.barrier()
logger.log("sampling complete")
def main():
args = create_argparser().parse_args()
dist_util.setup_dist()
logger.configure()
logger.log("creating model and diffusion...")
model, diffusion = create_model_and_diffusion(
**args_to_dict(args,
model_and_diffusion_defaults().keys()))
model.load_state_dict(
dist_util.load_state_dict(args.model_path, map_location="cpu"))
model.to(dist_util.dev())
if args.use_fp16:
model.convert_to_fp16()
model.eval()
logger.log("loading classifier...")
classifier = create_classifier(
**args_to_dict(args,
classifier_defaults().keys()))
classifier.load_state_dict(
dist_util.load_state_dict(args.classifier_path, map_location="cpu"))
classifier.to(dist_util.dev())
if args.classifier_use_fp16:
classifier.convert_to_fp16()
classifier.eval()
def cond_fn(x, t, y=None):
assert y is not None
with th.enable_grad():
x_in = x.detach().requires_grad_(True)
logits = classifier(x_in, t)
log_probs = F.log_softmax(logits, dim=-1)
selected = log_probs[range(len(logits)), y.view(-1)]
return th.autograd.grad(selected.sum(),
x_in)[0] * args.classifier_scale
def model_fn(x, t, y=None):
assert y is not None
return model(x, t, y if args.class_cond else None)
logger.log("sampling...")
all_images = []
all_labels = []
while len(all_images) * args.batch_size < args.num_samples:
model_kwargs = {}
classes = th.randint(low=0,
high=NUM_CLASSES,
size=(args.batch_size, ),
device=dist_util.dev())
model_kwargs["y"] = classes
sample_fn = (diffusion.p_sample_loop
if not args.use_ddim else diffusion.ddim_sample_loop)
sample = sample_fn(
model_fn,
(args.batch_size, 3, args.image_size, args.image_size),
clip_denoised=args.clip_denoised,
model_kwargs=model_kwargs,
cond_fn=cond_fn,
device=dist_util.dev(),
)
sample = ((sample + 1) * 127.5).clamp(0, 255).to(th.uint8)
sample = sample.permute(0, 2, 3, 1)
sample = sample.contiguous()
gathered_samples = [
th.zeros_like(sample) for _ in range(dist.get_world_size())
]
dist.all_gather(gathered_samples,
sample) # gather not supported with NCCL
all_images.extend(
[sample.cpu().numpy() for sample in gathered_samples])
gathered_labels = [
th.zeros_like(classes) for _ in range(dist.get_world_size())
]
dist.all_gather(gathered_labels, classes)
all_labels.extend([labels.cpu().numpy() for labels in gathered_labels])
logger.log(f"created {len(all_images) * args.batch_size} samples")
arr = np.concatenate(all_images, axis=0)
arr = arr[:args.num_samples]
label_arr = np.concatenate(all_labels, axis=0)
label_arr = label_arr[:args.num_samples]
if dist.get_rank() == 0:
shape_str = "x".join([str(x) for x in arr.shape])
out_path = os.path.join(logger.get_dir(), f"samples_{shape_str}.npz")
logger.log(f"saving to {out_path}")
np.savez(out_path, arr, label_arr)
dist.barrier()
logger.log("sampling complete")