U-Net by hand ✍️ ~ 17 steps walkthrough below
I consider U-Net as a key milestone in deep learning, the first image-to-image model that really worked!
It came out of medical imaging, an unusual place, not from NeurIPS or CVPR or ACL.
Now it is the backbone of diffusion models, which you see in almost all modern image generation models.
I drew the network as a C so the matrix multiplication flows naturally down.
Tilt your head to the right and it is a U again. 🤣
Goal: push a 3 x 16 image down to a 2 x 4 bottleneck and back out again, filling in every cell yourself.
= 1. Given =
An image of three channels, R, G and B, sixteen pixels wide, and every kernel the network will use.
= 2. Convolution 1 =
Let us slide the first kernel over the image. Each output is one multiply-and-add over a 2 x 3 window, and the result is the green feature map.
= 3. Find the maxima =
We circle the largest value in each 1 x 2 window. Circling first is worth the extra step: it is the pooling decision, made before anything is written down.
= 4. Max pool 1 =
Let us copy those maxima down. Sixteen columns become eight, and half the detail is gone for good.
= 5. Convolution 2 =
We convolve again with the second kernel, deeper into the contracting path. The feature map is blue now.
= 6. Find the maxima again =
Same move as step 3, on the blue map.
= 7. Max pool 2 =
Eight columns become four.
= 8. The bottleneck =
Let us convolve once more. This is the bottom of the U, a 2 x 4 block that is everything the network kept.
= 9. Spread it out =
We start back up. The transposed convolution writes each bottleneck value into a wider grid, leaving gaps between them.
= 10. Transposed convolution 1 =
Let us fill those gaps by convolving over the spread-out grid. Four columns become eight.
= 11. The first skip =
We copy the encoder's matching row straight across. This is the skip connection, and it is the whole reason a U-Net can recover detail that pooling threw away.
= 12. Convolution with the skip =
Let us convolve the upsampled features together with the copied ones.
= 13. Spread it out again =
Same as step 9, one level up.
= 14. Transposed convolution 2 =
Eight columns become sixteen, back to the width we started at.
= 15. The second skip =
The encoder's first feature map comes across, the one made before any pooling happened.
= 16. Convolution and ReLU =
We convolve, then cross out every negative and set it to zero.
= 17. Output convolution =
Let us apply the last kernel. Out comes R', G' and B', an image the same size as the one we started with.
The outputs:
R' = [3, 0, 7, 0, 7, 0, 17, 0, 3, 0, 9, 0, 2, 0, 6, 0]
G' = [1, 20, 1, 10, 1, 12, 1, 19, 2, 5, 1, 11, 1, 3, 1, 7]
B' = [4, 20, 8, 10, 8, 12, 18, 19, 5, 5, 10, 11, 3, 3, 7, 7]
Congrats! You just calculated a U-Net by hand.
💾 Save this post!
#UNet #DeepLearning #AI #NeuralNetworks #ComputerVision #MachineLearning
✨ Join Best TG Channels https://t.me/addlist/0f6vfFbEMdAwODBk
⭐️ Join Our WhatsApp Channel https://whatsapp.com/channel/0029VaC7Weq29753hpcggW2A
