lucidrains / big-sleep

A simple command line tool for text to image generation, using OpenAI's CLIP and a BigGAN. Technique was originally created by https://twitter.com/advadnoun
MIT License
2.57k stars 304 forks source link

Half Precision implemented for Low VRAM GPUs #127

Open cperales opened 2 years ago

cperales commented 2 years ago

Hi there!!

First of all, amazing job with your library. It achieves great results, and it is not difficult to implement.

I have several laptops, and I could run this and DeepDaze on my new laptop, but I couldn't run BigSleep on my old computer, with a GPU GTX 1050 with 4 GB of VRAM.

So I decided to implement a boolean parameter that reduces the precision of the model in the train_step method. Besides, I added another parameter, image_folder. This parameter can be a string, naming the folder where to save the images.

lucidrains commented 2 years ago

@cperales oh hey Carlos! this looks great! :heart: do you want to try extending this to the CLI as well?

cperales commented 2 years ago

Well, I am not so sure about how to do it... Is it enough if I modifed the file cli.py, adding the options to train function?

cperales commented 2 years ago

Btw, I found out that, better that some code can be simplified with context with torch.cuda.amp.autocast():.

https://spell.ml/blog/mixed-precision-training-with-pytorch-Xuk7YBEAACAASJam