U-Net from scratch

Hi,

I am working on implementing U-Net from scratch with the same architecture as it’s in the paper, guess I have built the model correct, but the problem is the model accuracy won’t increase or the DICE Score.

You can refer my code from my GitHub page, vence-andersen/UNet-PyTorch: A U-Net implementation from scratch. (github.com)

I am new to Deep Learning hence I am finding it difficult to troubleshoot this problem. Help is much appreciated, thanks in advance.

The model summary is below

==========================================================================================
Layer (type:depth-idx) Output Shape Param #

├─Sequential: 1-1 [-1, 64, 568, 568] –
| └─Conv2d: 2-1 [-1, 64, 570, 570] 1,792
| └─ReLU: 2-2 [-1, 64, 570, 570] –
| └─Conv2d: 2-3 [-1, 64, 568, 568] 36,928
| └─ReLU: 2-4 [-1, 64, 568, 568] –
├─MaxPool2d: 1-2 [-1, 64, 284, 284] –
├─Sequential: 1-3 [-1, 128, 280, 280] –
| └─Conv2d: 2-5 [-1, 128, 282, 282] 73,856
| └─ReLU: 2-6 [-1, 128, 282, 282] –
| └─Conv2d: 2-7 [-1, 128, 280, 280] 147,584
| └─ReLU: 2-8 [-1, 128, 280, 280] –
├─MaxPool2d: 1-4 [-1, 128, 140, 140] –
├─Sequential: 1-5 [-1, 256, 136, 136] –
| └─Conv2d: 2-9 [-1, 256, 138, 138] 295,168
| └─ReLU: 2-10 [-1, 256, 138, 138] –
| └─Conv2d: 2-11 [-1, 256, 136, 136] 590,080
| └─ReLU: 2-12 [-1, 256, 136, 136] –
├─MaxPool2d: 1-6 [-1, 256, 68, 68] –
├─Sequential: 1-7 [-1, 512, 64, 64] –
| └─Conv2d: 2-13 [-1, 512, 66, 66] 1,180,160
| └─ReLU: 2-14 [-1, 512, 66, 66] –
| └─Conv2d: 2-15 [-1, 512, 64, 64] 2,359,808
| └─ReLU: 2-16 [-1, 512, 64, 64] –
├─MaxPool2d: 1-8 [-1, 512, 32, 32] –
├─Sequential: 1-9 [-1, 1024, 28, 28] –
| └─Conv2d: 2-17 [-1, 1024, 30, 30] 4,719,616
| └─ReLU: 2-18 [-1, 1024, 30, 30] –
| └─Conv2d: 2-19 [-1, 1024, 28, 28] 9,438,208
| └─ReLU: 2-20 [-1, 1024, 28, 28] –
├─ConvTranspose2d: 1-10 [-1, 512, 56, 56] 2,097,664
├─Sequential: 1-11 [-1, 512, 52, 52] –
| └─Conv2d: 2-21 [-1, 512, 54, 54] 4,719,104
| └─ReLU: 2-22 [-1, 512, 54, 54] –
| └─Conv2d: 2-23 [-1, 512, 52, 52] 2,359,808
| └─ReLU: 2-24 [-1, 512, 52, 52] –
├─ConvTranspose2d: 1-12 [-1, 256, 104, 104] 524,544
├─Sequential: 1-13 [-1, 256, 100, 100] –
| └─Conv2d: 2-25 [-1, 256, 102, 102] 1,179,904
| └─ReLU: 2-26 [-1, 256, 102, 102] –
| └─Conv2d: 2-27 [-1, 256, 100, 100] 590,080
| └─ReLU: 2-28 [-1, 256, 100, 100] –
├─ConvTranspose2d: 1-14 [-1, 128, 200, 200] 131,200
├─Sequential: 1-15 [-1, 128, 196, 196] –
| └─Conv2d: 2-29 [-1, 128, 198, 198] 295,040
| └─ReLU: 2-30 [-1, 128, 198, 198] –
| └─Conv2d: 2-31 [-1, 128, 196, 196] 147,584
| └─ReLU: 2-32 [-1, 128, 196, 196] –
├─ConvTranspose2d: 1-16 [-1, 64, 392, 392] 32,832
├─Sequential: 1-17 [-1, 64, 388, 388] –
| └─Conv2d: 2-33 [-1, 64, 390, 390] 73,792
| └─ReLU: 2-34 [-1, 64, 390, 390] –
| └─Conv2d: 2-35 [-1, 64, 388, 388] 36,928
| └─ReLU: 2-36 [-1, 64, 388, 388] –
├─Conv2d: 1-18 [-1, 1, 388, 388] 65

Total params: 31,031,745
Trainable params: 31,031,745
Non-trainable params: 0
Total mult-adds (G): 167.72

Input size (MB): 3.74
Forward/backward pass size (MB): 1023.89
Params size (MB): 118.38
Estimated Total Size (MB): 1146.01