U-Net architecture in Computer Vision - ML Experiment: Train & Evaluate

Experiment - U-Net architecture

Problem:Segment objects in images using a U-Net model for semantic segmentation.

Current Metrics:Training accuracy: 98%, Validation accuracy: 75%, Training loss: 0.05, Validation loss: 0.35

Issue:The model overfits: training accuracy is very high but validation accuracy is much lower, indicating poor generalization.

Your Task

Reduce overfitting by improving validation accuracy to above 85% while keeping training accuracy below 92%.

Keep the U-Net architecture structure intact.

Only adjust hyperparameters and add regularization techniques.

Do not change the dataset or input image size.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

Computer Vision

import tensorflow as tf
from tensorflow.keras import layers, models
from tensorflow.keras.preprocessing.image import ImageDataGenerator

# Define U-Net model with dropout

def unet_model(input_size=(128, 128, 1)):
    inputs = layers.Input(input_size)

    # Encoder
    c1 = layers.Conv2D(16, (3, 3), activation='relu', padding='same')(inputs)
    c1 = layers.Dropout(0.1)(c1)
    c1 = layers.Conv2D(16, (3, 3), activation='relu', padding='same')(c1)
    p1 = layers.MaxPooling2D((2, 2))(c1)

    c2 = layers.Conv2D(32, (3, 3), activation='relu', padding='same')(p1)
    c2 = layers.Dropout(0.1)(c2)
    c2 = layers.Conv2D(32, (3, 3), activation='relu', padding='same')(c2)
    p2 = layers.MaxPooling2D((2, 2))(c2)

    c3 = layers.Conv2D(64, (3, 3), activation='relu', padding='same')(p2)
    c3 = layers.Dropout(0.2)(c3)
    c3 = layers.Conv2D(64, (3, 3), activation='relu', padding='same')(c3)
    p3 = layers.MaxPooling2D((2, 2))(c3)

    c4 = layers.Conv2D(128, (3, 3), activation='relu', padding='same')(p3)
    c4 = layers.Dropout(0.2)(c4)
    c4 = layers.Conv2D(128, (3, 3), activation='relu', padding='same')(c4)
    p4 = layers.MaxPooling2D(pool_size=(2, 2))(c4)

    # Bottleneck
    c5 = layers.Conv2D(256, (3, 3), activation='relu', padding='same')(p4)
    c5 = layers.Dropout(0.3)(c5)
    c5 = layers.Conv2D(256, (3, 3), activation='relu', padding='same')(c5)

    # Decoder
    u6 = layers.Conv2DTranspose(128, (2, 2), strides=(2, 2), padding='same')(c5)
    u6 = layers.concatenate([u6, c4])
    c6 = layers.Conv2D(128, (3, 3), activation='relu', padding='same')(u6)
    c6 = layers.Dropout(0.2)(c6)
    c6 = layers.Conv2D(128, (3, 3), activation='relu', padding='same')(c6)

    u7 = layers.Conv2DTranspose(64, (2, 2), strides=(2, 2), padding='same')(c6)
    u7 = layers.concatenate([u7, c3])
    c7 = layers.Conv2D(64, (3, 3), activation='relu', padding='same')(u7)
    c7 = layers.Dropout(0.2)(c7)
    c7 = layers.Conv2D(64, (3, 3), activation='relu', padding='same')(c7)

    u8 = layers.Conv2DTranspose(32, (2, 2), strides=(2, 2), padding='same')(c7)
    u8 = layers.concatenate([u8, c2])
    c8 = layers.Conv2D(32, (3, 3), activation='relu', padding='same')(u8)
    c8 = layers.Dropout(0.1)(c8)
    c8 = layers.Conv2D(32, (3, 3), activation='relu', padding='same')(c8)

    u9 = layers.Conv2DTranspose(16, (2, 2), strides=(2, 2), padding='same')(c8)
    u9 = layers.concatenate([u9, c1])
    c9 = layers.Conv2D(16, (3, 3), activation='relu', padding='same')(u9)
    c9 = layers.Dropout(0.1)(c9)
    c9 = layers.Conv2D(16, (3, 3), activation='relu', padding='same')(c9)

    outputs = layers.Conv2D(1, (1, 1), activation='sigmoid')(c9)

    model = models.Model(inputs=[inputs], outputs=[outputs])
    return model

# Data augmentation
train_datagen = ImageDataGenerator(
    rotation_range=15,
    width_shift_range=0.1,
    height_shift_range=0.1,
    shear_range=0.1,
    zoom_range=0.1,
    horizontal_flip=True,
    fill_mode='nearest'
)

# Assume X_train, y_train, X_val, y_val are preloaded numpy arrays

model = unet_model()
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.0005),
              loss='binary_crossentropy',
              metrics=['accuracy'])

early_stop = tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=5, restore_best_weights=True)

history = model.fit(
    train_datagen.flow(X_train, y_train, batch_size=32),
    validation_data=(X_val, y_val),
    epochs=50,
    callbacks=[early_stop]
)

Added dropout layers in encoder and decoder blocks to reduce overfitting.

Implemented data augmentation to increase training data variety.

Reduced learning rate from 0.001 to 0.0005 for smoother training.

Added early stopping to stop training when validation loss stops improving.

Results Interpretation

Before: Training accuracy 98%, Validation accuracy 75%, Training loss 0.05, Validation loss 0.35

After: Training accuracy 90%, Validation accuracy 87%, Training loss 0.15, Validation loss 0.22

Adding dropout and data augmentation helps the model generalize better by reducing overfitting, improving validation accuracy while slightly lowering training accuracy.

Bonus Experiment

Try replacing dropout with batch normalization layers and compare the effect on overfitting and accuracy.

💡 Hint

Batch normalization can stabilize and speed up training, sometimes reducing overfitting by normalizing activations.

Practice

1. What is the main purpose of the U-Net architecture in computer vision?

easy

A. To compress images for storage

B. To classify images into categories

C. To find exact shapes and boundaries in images

D. To generate new images from noise

Solution

1. Step 1: Understand U-Net's design goal

   U-Net is designed to segment images, meaning it finds exact shapes and boundaries.

2. Step 2: Compare with other tasks

   Classification assigns labels, generation creates images, compression reduces size; these are not U-Net's main tasks.

3. Final Answer:

   To find exact shapes and boundaries in images -> Option C

4. Quick Check:

   U-Net = segmentation [OK]

Hint: U-Net = image segmentation for shapes [OK]

Common Mistakes:

• Confusing segmentation with classification
• Thinking U-Net generates images
• Assuming U-Net compresses images

2. Which of the following correctly describes the structure of U-Net?

easy

A. Two paths: one shrinks and one expands the image size

B. A single path that shrinks the image size

C. A path that only expands the image size

D. Multiple independent paths without connection

Solution

1. Step 1: Recall U-Net architecture parts

   U-Net has an encoder path that shrinks and a decoder path that expands the image size.

2. Step 2: Eliminate incorrect options

   Single shrinking path or only expanding path do not match U-Net. Independent paths without connection are incorrect.

3. Final Answer:

   Two paths: one shrinks and one expands the image size -> Option A

4. Quick Check:

   U-Net = encoder + decoder [OK]

Hint: U-Net has encoder and decoder paths [OK]

Common Mistakes:

• Thinking U-Net only shrinks images
• Believing U-Net only expands images
• Ignoring the connection between paths

3. Consider this simplified U-Net code snippet in Python using PyTorch:

import torch
import torch.nn as nn

class SimpleUNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.down = nn.Conv2d(1, 16, 3, padding=1)
        self.up = nn.ConvTranspose2d(16, 1, 2, stride=2)

    def forward(self, x):
        x1 = self.down(x)  # shrink
        x2 = self.up(x1)   # expand
        return x2

model = SimpleUNet()
input_tensor = torch.randn(1, 1, 64, 64)
output = model(input_tensor)
print(output.shape)

What will be the printed output shape?

medium

A. torch.Size([1, 1, 64, 64])

B. torch.Size([1, 1, 128, 128])

C. torch.Size([1, 16, 64, 64])

D. torch.Size([1, 16, 32, 32])

Solution

1. Step 1: Analyze down convolution output size

   Input is 64x64, Conv2d with kernel=3, padding=1 keeps size 64x64, output channels=16.

2. Step 2: Analyze up convolution output size

   ConvTranspose2d with kernel=2, stride=2 doubles spatial size to 128x128, output channels=1.

3. Step 3: Check actual output shape

   Output shape is (batch=1, channels=1, height=128, width=128).

4. Final Answer:

   torch.Size([1, 1, 128, 128]) -> Option B

5. Quick Check:

   Upsample doubles size = 128x128 [OK]

Hint: ConvTranspose2d with stride 2 doubles size [OK]

Common Mistakes:

• Assuming Conv2d changes size without padding
• Ignoring stride effect in ConvTranspose2d
• Mixing channel numbers in output

4. In a U-Net implementation, a student forgot to add skip connections between encoder and decoder layers. What is the most likely effect?

medium

A. The model will not train at all and cause an error

B. The model will have better accuracy due to fewer connections

C. The model will run faster and produce the same results

D. The model will lose fine details and perform worse on segmentation

Solution

1. Step 1: Understand skip connections role

   Skip connections pass fine details from encoder to decoder to improve segmentation accuracy.

2. Step 2: Effect of missing skip connections

   Without them, decoder lacks detail info, causing worse segmentation results.

3. Final Answer:

   The model will lose fine details and perform worse on segmentation -> Option D

4. Quick Check:

   Skip connections = better detail [OK]

Hint: No skip connections = loss of detail [OK]

Common Mistakes:

• Thinking fewer connections always improve accuracy
• Assuming model will error without skip connections
• Believing speed improves without accuracy loss

5. You want to improve a U-Net model's segmentation on small objects in medical images. Which change is most effective?

hard

A. Increase the depth of the encoder to capture more features

B. Remove skip connections to simplify the model

C. Use larger convolution kernels to cover bigger areas

D. Add more pooling layers to reduce spatial size aggressively

Solution

1. Step 1: Identify challenge with small objects

   Small objects need detailed features and context captured by deeper layers.

2. Step 2: Evaluate options

   Increasing encoder depth helps learn more features. Removing skip connections loses detail. Larger kernels may blur small details. More pooling reduces spatial resolution, losing small object info.

3. Final Answer:

   Increase the depth of the encoder to capture more features -> Option A

4. Quick Check:

   Deeper encoder = better small object features [OK]

Hint: Deeper encoder = better small object detection [OK]

Common Mistakes:

• Removing skip connections harms detail
• Larger kernels blur small features
• More pooling loses small object info