在 Keras 中实现注意力 r2unet 网络 - tf 2.x

Question

我想使用注意力网络从眼底图像中分割视盘。该模型的架构选自“https://github.com/lixiaolei1982/Keras-Implementation-of-U-Net-R2U-Net-Attention-U-Net-Attention-R2U-Net.-”（礼貌，归功于 - lixiaolei1982）。 当我跟踪训练和验证损失时，它会减少（见下图），但训练损失在 20 个 epoch 后接近 0。在前 20 个 epoch 中，分割后的图像是全黑或全白的。我试图标准化预测的图像。但它是一样的。 谁能帮我纠正这个问题？是损失函数的作用导致输出图像全黑吗？下面是使用注意力网络训练图像的代码

import os
import pandas as pd
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import numpy as np 
import matplotlib.pyplot as plt
from tensorflow.keras.backend import flatten
from skimage.exposure import equalize_hist as clhe
import network as new_model

batch_size = 32
no_epochs = 20
img_height, img_width, img_num_channels = 512, 512, 3

mdl1 = "att_r2unet"

m1 = new_model.att_r2_unet(img_height, img_width, n_label=1)

def preprocess_im(imgs):
    """Make input image values lie between -1 and 1."""
    #imgs = clhe(imgs)
    out_imgs = imgs - np.max(imgs)/2.
    out_imgs /= np.max(imgs)/2.
    return out_imgs

##

seg_train_gen = ImageDataGenerator(preprocessing_function=preprocess_im)
seg_train_x = seg_train_gen.flow_from_directory(directory="../train_im",target_size=(512,512),batch_size=batch_size,color_mode="rgb",class_mode="sparse",shuffle=True,seed=30)
seg_train_y = seg_train_gen.flow_from_directory(directory="../train_gt",target_size=(512,512),batch_size=batch_size,color_mode="grayscale",class_mode="sparse",shuffle=True,seed=30)
seg_val_x = seg_train_gen.flow_from_directory(directory="../val_im",target_size=(512,512),batch_size=batch_size,color_mode="rgb",class_mode="sparse",shuffle=True,seed=30)
seg_val_y = seg_train_gen.flow_from_directory(directory="../val_gt",target_size=(512,512),batch_size=batch_size,color_mode="grayscale",class_mode="sparse",shuffle=True,seed=30)

b_iter = int(np.ceil(seg_train_x.n / batch_size))

hist = {'m1_loss': [],'m1_vloss': []}

epoch = 0
for e in range(epoch, no_epochs):
l1 = list(np.zeros(b_iter))  
    for it in range(b_iter):
        x_batch,_ = seg_train_x.next() 
        y_batch, _ = seg_train_y.next()
        loss1 = m1.train_on_batch(x_batch, y_batch)
        l1.append(loss1[0])
    xval,_ = seg_val_x.next()
    yval,_ = seg_val_y.next()
    loss2a = m1.evaluate(xval,yval, batch_size=batch_size, verbose=1) 
    print('Epoch %d / %d tr_loss %.6f val_loss %.6f, ' % (e + 1, no_epochs, np.mean(l1), loss2a[0]))


    hist['m1_vloss'].append(loss2a[0])
 
    # Save best model
    if e > epoch+1:

        Eopt1 = np.min(hist['m1_vloss'][:-1])

        if hist['m1_vloss'][-1] < Eopt1:
            m1.save((save_dir+mdl1+'_best_model.h5'),overwrite=True)
            m1.save_weights((save_dir+mdl1+'_best_weights.h5'),overwrite=True)

    # save intermediate to folder results every 2 epochs
    if e % 2 == 0:

        x_plt = (xval[0] - xval[0].min()) / (xval[0].max() - xval[0].min())

        ypred = m1.predict(xval)
        fix, ax = plt.subplots(1,3, figsize=(10,10))
        ax[0].imshow(x_plt)
        ax[1].imshow(yval[0, :, :, 0], cmap='gray')
        ax[2].imshow(ypred[0, :, :, 0], cmap='gray')
        plt.savefig((save_dir+mdl1+'e_' + str(e) + '.jpg'))
        plt.close()
  
print("Completed training...")

我已经尝试添加批量标准化，并且还在最终层中将激活函数从 sigmoid 修改为 relu。也尝试更改优化器。但这些都没有帮助。事实上，有时不修改架构，损失是负数。

下面是训练时保存的样本输出（左边是输入图像，中间是ground truth，右边是预测输出） 感谢您的宝贵时间。

Answer 1

您可以检查学习率，使用的优化器类型。如果结果图像是全白或全黑，则可能是您的数据集类别不平衡。

某个班级压倒另一个班级。您必须平衡数据集或使用更好的损失函数，如 iou、dice 或焦点损失。