How to deal with unbalanced data in pixelwise classification?

Question

I'm trying to train a fully convolutional network to generate binary image showing roads from satellite photos. For a long time, I struggled to get the network to output anything but black - then I figured out the problem. I have plenty of training data, but of course most of the output pixels will be black, since roads take only small percentage of space.

I've created a toy example showing the problem:

from keras.models import Sequential
from keras.layers import Dense, Flatten, Conv2D, MaxPooling2D
import numpy as np
import random
import sklearn.model_selection

THRESHOLD = 3
BALANCED = False

if BALANCED:
    X = [[
            random.randint(THRESHOLD + 1, 255) if
                random.randint(0, 1) else
            random.randint(0, THRESHOLD)
            for _ in range(10)] for _ in range(10**4)]
    Y = [[0] if x[0] <= THRESHOLD else [1] for x in X]
else:
    X = [[random.randint(0, 255) for _ in range(10)] for _ in range(10**4)]
    Y = [[0] if x[0] <= THRESHOLD else [1] for x in X]

X, Y = np.array(X), np.array(Y)
X_train, X_test, Y_train, Y_test = sklearn.model_selection.train_test_split(
        X, Y, test_size = 0.2)

model = Sequential()
model.add(Dense(8, activation = "sigmoid", input_shape = (10, )))
model.add(Dense(1, activation = "sigmoid"))

model.compile(loss = "binary_crossentropy", optimizer = "adam")

model.summary()
model.fit(X_train, Y_train, 
        batch_size = 10, epochs = 10, 
        verbose = 2, 
        validation_data = (X_test, Y_test))

Y_res = model.predict(X_test)
Ys0 = [y[0] for x, y, y_opt in zip(X_test, Y_res, Y_test) if y_opt[0] == 0]
Ys1 = [y[0] for x, y, y_opt in zip(X_test, Y_res, Y_test) if y_opt[0] == 1]

print(np.mean(Ys0))
print(np.mean(Ys1))

Here, I artificially create data, with X being ten random bytes, and Y being boolean thresholding first byte on the value of three. That means normally only around 1/64 of all samples will be generated with y equal to 0, similarly to my original problem. There's also a second dataset, with the same decision boundary, but more evenly (class-wise) distributed data - set by changing BALANCED toggle in the code. The results confirm my suspicions.

For balanced set, actual 0's are given, on average, value of 0.025. 1's - 0.97. That's reasonable. On the unbalanced, raw dataset, 0's have mean of 0.87, and 1's - 0.98. That means the zeros are pretty much always classified as one anyway. On my real data, the effect is even worse, and the returned value is pretty much equal to percentage of black, regardless of real classification.

I tried adding class_weights = [100, 1] to fit call, but it didn't change much. Is there a good strategy for dealing with such unbalanced data? In this toy example I could probably sample the x-y pairs (e.g. keep only 5% of 1-classified data), but I don't think this would work for my full problem - since I'm using FCN, I'm putting a whole image into the network at once. There are no images with 50% roads, so I cannot possibly sample images to make roads take enough space.

I also tried, quite randomly, changing loss functions, including using one for categorical data and one-hot encoding expected outputs (i.e. [1, 0] and [0, 1] instead of 0 and 1). This did not fix anything either.

Answer 1

Looks like the problem lied just in poorly chosen parameters. In the example given above, simply increasing epoch count significantly improved the results. In my actual image processing problem, I had to decrease learning rate, since with the default one, ReLU's were saturating very quickly, hindering further progress.