diff --git a/MindIE/MindIE-Torch/built-in/foundation/DiT/README.md b/MindIE/MindIE-Torch/built-in/foundation/DiT/README.md index c13e1e3a0c476aa472196f005cf13ea72c15cc9d..2cd50ecea7fb59d54930b345a77f89f7f898476c 100644 --- a/MindIE/MindIE-Torch/built-in/foundation/DiT/README.md +++ b/MindIE/MindIE-Torch/built-in/foundation/DiT/README.md @@ -47,7 +47,7 @@ latent_size = image_size // 8 ```bash git clone https://github.com/facebookresearch/DiT - mv background_runtime.py export_model.py models_npu.py sample_npu.py timm_patch.py requirements.txt ./DiT + mv background_runtime.py export_model.py models_npu.py sample_npu.py vision.patch timm_patch.py requirements.txt ./DiT ``` 2. 安装依赖 @@ -120,7 +120,7 @@ latent_size = image_size // 8 - --ckpt：DiT-XL-2的权重路径 - --vae_model：vae的权重路径 - - --image_size：分辨率，默认为512 + - --image_size：分辨率，支持256和512。默认为512 - --device：使用哪张卡 - --soc：soc_version，只支持Duo和A2 - --output_dir：pt模型输出目录 @@ -162,7 +162,7 @@ latent_size = image_size // 8 参数说明： - --vae：使用哪种vae模型，支持mse和ema - - --image_size：分辨率，默认为512 + - --image_size：分辨率，支持256和512。默认为512 - --ckpt：DiT-XL-2的权重路径 - --device：使用哪张卡 - --class_label：可在0~999中任意指定一个整数，代表image_net的种类 @@ -171,7 +171,7 @@ latent_size = image_size // 8 4. 精度验证 - 下载数据集[ImageNet512x512](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/VIRTUAL_imagenet512.npz)，放在任意路径 + 下载数据集[ImageNet512x512](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/VIRTUAL_imagenet512.npz)和[ImageNet256x256](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/VIRTUAL_imagenet256_labeled.npz)，放在任意路径 然后执行以下命令： @@ -187,7 +187,7 @@ latent_size = image_size // 8 --results results # A2 - python3 sample_npu.py \ + python3 fid_test.py \ --vae mse \ --image_size 512 \ --ckpt ./DiT-XL-2-512x512.pt \ @@ -199,9 +199,10 @@ latent_size = image_size // 8 参数说明： - --results：生成的1000张图片存放路径 - + - image_size：分辨率，支持256和512。默认为512 + 之后进行FID计算： - + ```bash python3 -m pytorch_fid ./VIRTUAL_imagenet512.npz ./results ``` @@ -210,7 +211,9 @@ latent_size = image_size // 8 性能参考下列数据。 -| 硬件形态 | 迭代次数 | 平均耗时 | -| -------- | -------- | -------- | -| Duo | 250 | 19.6s | -| A2 | 250 | 10.49s | +| 分辨率 | 硬件形态 | 迭代次数 | 平均耗时 | +| ------ | -------- | -------- | -------- | +| 512 | Duo | 250 | 19.6s | +| | A2 | 250 | 10.49s | +| 256 | Duo | 250 | 9.5s | +| | A2 | 50 | 4.13s | diff --git a/MindIE/MindIE-Torch/built-in/foundation/DiT/fid_test.py b/MindIE/MindIE-Torch/built-in/foundation/DiT/fid_test.py index 2f4e3f6faaabbb18eaf0f7332c723d4e4b1d8e67..5b360f7a233c62e32b052169ff4f3e58bc5edfd9 100644 --- a/MindIE/MindIE-Torch/built-in/foundation/DiT/fid_test.py +++ b/MindIE/MindIE-Torch/built-in/foundation/DiT/fid_test.py @@ -53,6 +53,7 @@ def main(args): mindie_model_path = f"{args.output_dir}/dit/dit_model_{args.image_size}_compiled.pt" vae_compiled_model_path = f"{args.output_dir}/vae/vae_{args.vae}_{args.image_size}_compiled.pt" vae_compiled_model = torch.jit.load(vae_compiled_model_path).eval() + dit_compiled_model = torch.jit.load(mindie_model_path).eval() if args.ckpt is None: assert args.model == "DiT-XL/2", "Only DiT-XL/2 models are available for auto-download." @@ -72,7 +73,7 @@ def main(args): diffusion = create_diffusion(str(args.num_sampling_steps)) mindietorch.set_device(device_id) - model.set_npu_model_stream(args.parallel, device_id, args.image_size, mindie_model_path) + model.set_npu_model_stream(args.parallel, device_id, args.image_size, mindie_model_path, dit_compiled_model) all_class = 1000 for i in range(all_class): diff --git a/MindIE/MindIE-Torch/built-in/foundation/DiT/models_npu.py b/MindIE/MindIE-Torch/built-in/foundation/DiT/models_npu.py index 1ca5f528f155471b7bc3df793f94c0e230adb700..796fcc9a47ba7da94670fa16254b07ca744c0b7c 100644 --- a/MindIE/MindIE-Torch/built-in/foundation/DiT/models_npu.py +++ b/MindIE/MindIE-Torch/built-in/foundation/DiT/models_npu.py @@ -35,12 +35,12 @@ class MindIEDiT(DiT): eps = torch.cat([half_eps, half_eps], dim=0) return torch.cat([eps, rest], dim=1) - def set_npu_model_stream(self, parallel, device, image_size, mindie_model_path): + def set_npu_model_stream(self, parallel, device, image_size, mindie_model_path, dit_compiled_model): latent_size = image_size // 8 self.device, device_2 = device, device + 1 self.stream = mindietorch.npu.Stream(f"npu:{self.device}") self.parallel = parallel - self.model_npu = torch.jit.load(mindie_model_path).eval() + self.model_npu = dit_compiled_model if parallel: runtime_info = RuntimeIOInfo( input_shapes=[(1, 4, latent_size, latent_size), (1,), (1,)], diff --git a/MindIE/MindIE-Torch/built-in/foundation/DiT/sample_npu.py b/MindIE/MindIE-Torch/built-in/foundation/DiT/sample_npu.py index bbd4c7502c8faf78e4621c54dc4dadf4764d091e..6595e3395499849654569ffd312eede888b95f7a 100644 --- a/MindIE/MindIE-Torch/built-in/foundation/DiT/sample_npu.py +++ b/MindIE/MindIE-Torch/built-in/foundation/DiT/sample_npu.py @@ -46,6 +46,7 @@ def main(args): mindie_model_path = f"{args.output_dir}/dit/dit_model_{args.image_size}_compiled.pt" vae_compiled_model_path = f"{args.output_dir}/vae/vae_{args.vae}_{args.image_size}_compiled.pt" vae_compiled_model = torch.jit.load(vae_compiled_model_path).eval() + dit_compiled_model = torch.jit.load(mindie_model_path).eval() if args.ckpt is None: assert args.model == "DiT-XL/2", "Only DiT-XL/2 models are available for auto-download." @@ -77,7 +78,7 @@ def main(args): model_kwargs = dict(y=y, cfg_scale=args.cfg_scale) mindietorch.set_device(device_id) - model.set_npu_model_stream(args.parallel, device_id, args.image_size, mindie_model_path) + model.set_npu_model_stream(args.parallel, device_id, args.image_size, mindie_model_path, dit_compiled_model) start = time.time() samples = diffusion.p_sample_loop( model.forward_with_cfg,