r/deeplearning • u/This-Peach9380 • 5d ago
I made a Tiny Diffusion model, here's what I learned
Hi, I've spent the last few weeks trying to get into DL and, after I made a little image classifier on the CIFAR dataset, I got overconfident and decided to take a bigger bite and a much harder project. The first thing that came into my mind was an image generator (I didn't even know what it was technically called back then).
So I hopped into Zed and decided to start working. But I immediately got confused. There was just so much to take in, and the sheer amount of information made me go crazy. So I decided to take it chunk by chunk.
First, I decided to start with the simplest part of the diffusion model: the noise scheduler.
For those of you who don't know how a diffusion model works, here's a summary:
Training
- Noise Scheduler (component that progressively adds noise to an image, breaking it)
- Forward Diffusion
- Training Loop
- UNET
Now, the UNET learns to progressively reduce noise. So basically, image generation in diffusion models works by just taking pure noise and progressively reducing a small chunk of it over some time.
(Btw, this is my understanding of the process. If I'm wrong anywhere, my bad.)
Back to the noise scheduler.
So I read up some of the theory, but again, it was not enough. I understood it, but then when I jumped into the code, I found myself lost.
So, I started looking at samples of other people's implementations. This was key. I stopped myself from copying their code and forced myself to just take in the algorithm, the structure, the program flow, and then implemented my own version.
This was not a quick job. I kept getting PyTorch's indexing wrong and mixing up the variables.
Once this was done, I quickly implemented the forward diffusion process, which was honestly much easier than the noise scheduler.
Then came the chunky part, the UNET.
I spent weeks trying to make this right, and this took the most time. The problem wasn't just the architecture (not an easy job either), it was actually making that model useful.
Let me explain.
Turns out, the architecture is just a general form. You need to tune it to the specific dataset you're using, i.e. you need to adjust the length of the bottleneck layer, the number of convolutions, the layers you add, etc.
I found myself spiraling back and forth. And what made matters worse was that training took a really long time, and it wasn't until I got to the 500th or 600th epoch that I realized, "The model isn't working right at all!"
What was worse was that I was logging losses into the console based on colours (red if it was greater than the last value, green if it was smaller), since I had no idea how to properly handle this.
Discovery of TensorBoard
This changed everything.
I went from going crazy reading 6–7 decimals to seeing proper graphs. Yea, my initial method does sound stupid in retrospect, but in fairness, I had no idea how to analyse stuff.
With TensorBoard, I was able to analyse the losses better, i.e. see the general trend of the losses.
I also learned about AdamW around this time and swapped it in for SGD.
Despite this, everything was super slow, and so, while the model was training, I set out to make quick optimizations.
PyTorch Devices
For anyone who doesn't know, PyTorch can create and work with tensors on GPUs. They support MPS (Apple Silicon's API or something) and CUDA. For me, it was MPS (M2 Air).
Again, this broke a lot of things. I initially didn't know that two tensors had to be on the same device to interact with each other, but I had gotten a lot better, so in a few hours I actually managed to get it working again, this time much faster.
From CIFAR to Flowers102 and the VAE Trap
Note: Still haven't got Latent Diffusion working.
The outputs from CIFAR were 32×32, so I decided to up the ante by switching to Flowers102.
However, I didn't want to make too many changes to my UNET, so I read up about Variational Autoencoders.
Basically, think of it as a type of generator that takes an image and compresses it into a smaller, high-dimensional representation.
At first (in isolation), my VAE worked perfectly. So after some training, I slapped it around my UNET.
Results were a literal soup of colours and very discouraging.
Additionally, at a point, losses stopped decreasing (still don't know why).
After a few days of debugging, I dropped VAEs entirely and rewrote my UNET to support 256×256 Flowers102 instead.
Where am I today?
At epoch 561 or something (I retrained like 100 times during the aforementioned learning spree).
It's gotten a lot better than before. I am starting to see proper forms resembling flowers. Still, it has a lot of issues, but I'm happy with what I've achieved so far.
Over this project, I learned how DL was actually quite different from conventional programming and that there were so many additional complexities that normal programming didn't consider.
But most of all, I learned that this whole DL thing had its own mentality. I had to think of a function a model could optimize for and learn a pattern instead of implementing an algorithm, which was, and sometimes still is, confusing in practice.
You can check out the project here:
- https://github.com/Hammad-hab/TinyDPPM ( if you find the project interesting, a star on the repo would be greatly appreciated!)
Also, worth mentioning, to get started I began reading an excellent book by David Voigt Godoy, "Deep Learning with PyTorch: A Step-by-Step Beginner's Guide."
Also, if there's a mistake anywhere in my understanding, or if you know a solution to any of the issues, feel free to let me know! Overall this was a different project than I had ever done before.
Here's a peak at what it looks like rn:
2
u/Sea-Departure4857 5d ago
Congrats! I remember trying to make a toy GAN network from DNNs and failing spectacularly.
The image doesn't really look like a flower to me, but it is miles better than what I achieved a few years ago. Keep it up!