stablediffusion icon indicating copy to clipboard operation
stablediffusion copied to clipboard

Switch from no_grad to inference_mode for scripts

Open kjerk opened this issue 2 years ago • 0 comments

Problem

torch.no_grad() is a slightly older API, pytorch advocates newer api for inference and future changes.

Proposed PR

  • torch.inference_mode() is a newer method which also disables grad, tracking, and a few other things for slightly better performance and is a transparent change.
  • This pr sweeps the scripts/ dir and replaces the no_grad() calls with inference_mode() context manager instead.
  • This is the advocated method in performance docs and by Pytorch themselves. The performance gain here for this model is negligible but it's best practice and more idiomatic for inference.

Testing Method

  • Main test: simple command repeatedly on main and dev branch python -m scripts.txt2img --prompt "a beautiful painting of an astronaut riding a horse" --ckpt /path/to/512-base-ema.ckpt --W 512 --H 512 --n_iter 4 --n_sample 2
  • The results are effectively identical within margin of error. Showing possibly an extremely tiny gain over time for the diffusion process though not always. Likely perf noise.

Before-After

before_after

System info

  • Windows 10 21H2
  • torch 1.12.1+cu116
  • xformers installed and working
  • cuda_11.6.r11.6/compiler.31057947_0

kjerk avatar Nov 24 '22 11:11 kjerk