Получение более точных результатов сегментации Deeplab для небольших / несбалансированных классов

Я обучил официальную модель deeplab Tensorflow (https://github.com/tensorflow/models/tree/master/research/deeplab) в наборе данных CelebAMask-HQ (https://github.com/switchablenorms/CelebAMask-HQ), чтобы иметь модель, которая может семантически сегментировать все лицевые сегменты (например, глаза, нос и т. д.). Я потратил немного времени на настройки гиперпараметров и использую настройки по умолчанию:

CUDA_VISIBLE_DEVICES=0 python "${WORK_DIR}"/deeplab/train.py \
  --logtostderr \
  --train_split="train" \
  --model_variant="xception_65" \
  --atrous_rates=6 \
  --atrous_rates=12 \
  --atrous_rates=18 \
  --output_stride=16 \
  --decoder_output_stride=4 \
  --train_crop_size="1025,1025" \
  --train_batch_size=2 \
  --training_number_of_steps=45000 \
  --fine_tune_batch_norm=false \
  --tf_initial_checkpoint="${WORK_DIR}/deeplab/pretrained_models/deeplabv3_pascal_trainval/model.ckpt" \
  --train_logdir="${WORK_DIR}/deeplab/logs" \
  --dataset="celeba" \
  --dataset_dir="${WORK_DIR}/deeplab/datasets/celeba/tfrecords_padded/"

Единственное, что я адаптировал, - это веса классов, в которых я вычислил веса классов на основе соотношения всех пикселей, принадлежащих каждому классу в общем наборе данных. Эти расчетные отношения следующие:

class_ratio = {0: 0.287781127731224, #BG
                       1: 0.31428004829848194, #hair
                       2: 0.25334614328648697, #face
                       3: 0.008209905199792278, #brows
                       4: 0.0044636011242926155, #eyes
                       5: 0.020564768086557928, #nose
                       6: 0.004150659950132944, #u_lip
                       7: 0.00680743101856918, #l_lip
                       8: 0.0030163743167156494, #mouth
                       9: 0.040800302545885576, #neck
                       10: 0.008106960279456135, #ears
                       11: 0.03355246488702522, #clothes
                       12: 0.009293231642880359, #hat
                       13: 0, #ear_ring -> 0
                       14: 0, #glasses -> 0
                       15: 0 #necklace -> 0
                       }

В качестве веса класса я беру 1/<class_ratio>, поэтому вес класса для фона равен 3,57, а для бровей - 121,95.

Наконец, я увеличиваю данные, например вращаю, переворачиваю и изменяю яркость.

Мои результаты довольно хороши, но есть кое-что заметное, когда я ввожу в модель некоторые из изображений своего тренировочного набора. Ниже исходной сегментации:

И вот сегментированный результат:

Как видите, сегментированный результат неплох, но особенно мелкие классы, такие как глаза, брови и нос, не так сильно сегментированы, как мне бы хотелось. В основном для всех изображений, сегментированных моделью, глаза, нос и брови больше, чем исходная сегментация. Поэтому я хотел бы изменить несколько гиперпараметров, чтобы получить более точные результаты сегментации для меньших классов.

Есть ли какие-либо предложения по возможному подходу к достижению более высоких результатов в мышцах меньших классов? Мой текущий подход к вычислению веса класса на основе абсолютного процента пикселей, принадлежащих каждому классу в общем наборе данных, работает достаточно хорошо, но, может быть, альтернативный подход для расчета весов классов работает лучше? Или другая базовая структура модели, способная выполнять точную сегментацию?

Любая помощь приветствуется. Спасибо!