| 81 | noise_scheduler = DDPMScheduler.from_pretrained("ashawkey/imagedream-ipmv-diffusers", subfolder="scheduler", revision=None) |
| 82 | |
| 83 | def print_model_info(model): |
| 84 | print("="*20) |
| 85 | print("model name: ", type(model).__name__) |
| 86 | print("learnable parameters(M): ", sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6) |
| 87 | print("non-learnable parameters(M): ", sum(p.numel() for p in model.parameters() if not p.requires_grad) / 1e6) |
| 88 | print("total parameters(M): ", sum(p.numel() for p in model.parameters()) / 1e6) |
| 89 | print("model size(MB): ", sum(p.numel() * p.element_size() for p in model.parameters()) / 1024 / 1024) |
| 90 | |
| 91 | if accelerator.is_main_process: |
| 92 | tracker_config = dict(vars(opt)) |