qubvel
can we add dropout to Unet
Is there a way to add dropout to the Unet model here? I see there is some dropout in FPN and PSPNet, but none in Unet. I want to train on some limited data and I'm having problems with overfitting, so some dropout could really help if there's a good place to insert dropout layers into this model.
I had a similar question recently and did some experimentation. I didn't get much improvement in my model, but maybe yours will benefit more.
In my branch of segmentation_models, unet takes an additional argument (center_dropout=0.0) which specifies how much dropout to apply at the center of the model, between the encoder and decoder.
I also experimented with adding dropout to the skip connections. This seemed to decrease accuracy a lot so it's commented out. But the line remains in my code and you can un-comment it if you want.
I added one dropout before the last activation.
Sample code,
from tensorflow.keras import layers
from tensorflow import keras
import segmentation_models as sm
def get_model_with_dropout(base_model_name='efficientnetb4',activation='sigmoid',dropout = 0.1):
base_model = sm.Unet(base_model_name, encoder_weights='imagenet')
base_model_input = base_model.input
base_model_output = base_model.get_layer('final_conv').output
#add dropout
base_model_output = keras.layers.Dropout(dropout)(base_model_output)
#add activation
output = keras.layers.Activation(activation, name=activation)(base_model_output)
model_dp = keras.models.Model(base_model_input, output)
return model_dp
hello @dennywangtenk did adding dropout help you overcome the problem of overfitting ?
