[Reproduceability 1/3] Allow tensors to be generated on CPU #1902
Conversation
|
The documentation is not available anymore as the PR was closed or merged. |
|
|
||
|
|
||
| # <original>.time_embed -> <diffusers>.time_embedding |
There was a problem hiding this comment.
Yeah, updated my black and then re-updated back to 22.8. We should maybe soon blackify the complete codebase once :-)
Co-authored-by: Anton Lozhkov <[email protected]>
pcuenca
left a comment
pcuenca
left a comment
There was a problem hiding this comment.
Haven't tested it yet but it looks great!
My only small concern is about the name, wrote a comment below.
| latents = torch.randn(shape, generator=generator, device="cpu", dtype=dtype).to(device) | ||
| else: | ||
| latents = torch.randn(shape, generator=generator, device=device, dtype=dtype) | ||
| latents = torch_randn(shape, generator=generator, device=device, dtype=dtype) |
| latents = [ | ||
| torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype) for i in range(batch_size) | ||
| ] | ||
| latents = torch.cat(latents, dim=0).to(device) |
There was a problem hiding this comment.
Very cool to include per-item reproducibility too
| logger = logging.get_logger(__name__) # pylint: disable=invalid-name | ||
|
|
||
|
|
||
| def torch_randn( |
There was a problem hiding this comment.
My only concern here is that torch_randn is easy to confuse (both visually and inadvertently while typing) with torch.randn. Would it make sense to make the name slightly more different? Can't think of anything great though, diffusers_randn feels kind of ugly.
There was a problem hiding this comment.
Maybe just randn_tensor?
Co-authored-by: Pedro Cuenca <[email protected]>
| f" Tensors will be created on 'cpu' and then moved to {device}. Note that one can probably" | ||
| f" slighly speed up this function by passing a generator that was created on the {device} device." | ||
| ) | ||
| elif generator.device.type != device.type and generator.device.type == "cuda": |
There was a problem hiding this comment.
I would add a comment here that we only allow cpu->cuda generation for reproducibility reasons, which is why the other way around is not supported / doesn't make sense.
There was a problem hiding this comment.
Yes, I'll make a whole doc page about this in the follow-up PR :-)
patrickvonplaten
changed the title
…gingface#1902) * [Deterministic torch randn] Allow tensors to be generated on CPU * fix more * up * fix more * up * Update src/diffusers/utils/torch_utils.py Co-authored-by: Anton Lozhkov <[email protected]> * Apply suggestions from code review * up * up * Apply suggestions from code review Co-authored-by: Pedro Cuenca <[email protected]> Co-authored-by: Anton Lozhkov <[email protected]> Co-authored-by: Pedro Cuenca <[email protected]>
This PR adds a new helper function
torch_randnthat has three purposes:This is the first PR of a two PR series to see if this helps to make UnCLIP deterministic across different CUDA, GPUs. In a follow-up PR, I'd then like to replace all existing
torch.randnfunctions with this function as well as add areproducibility.mdguide.