基于segmentation_models与Keras的多GPU训练异常问题排查

背景 #

单GPU（12GB显存）训练时，因显存限制只能设置batch size=1，此时模型训练正常，验证IOU可达0.5-0.6。为解决显存瓶颈，参考Keras多GPU文档，基于segmentation_models实现MirroredStrategy多GPU训练，但训练结果异常：IOU曲线不符合预期，快速趋于平稳或出现异常波动。测试2GPU+batch size=8的组合，问题仍存在。

训练代码 #

def run_training(BACKBONE, n_classes, activation, metrics, EPOCHS, X_train, X_test, y_train_cat, y_test_cat):
    # Create a MirroredStrategy.
    strategy = tf.distribute.MirroredStrategy()
    print('Number of devices: {}'.format(strategy.num_replicas_in_sync))

    # Open a strategy scope and create/restore the model
    with strategy.scope():
        model = sm.Unet(BACKBONE, encoder_weights='imagenet', classes=n_classes, activation=activation)
        weights = [0,.1,.45,.45]
        loss = weighted_categorical_crossentropy(weights)
        
        model.compile(optimizer='adam', loss=loss, metrics=metrics)
    #model1.compile(tf.keras.optimizers.legacy.Adam(), loss=loss, metrics=metrics)

        history=model.fit(X_train, 
          y_train_cat,
          batch_size=4, 
          epochs=EPOCHS,
          verbose=1,
          validation_data=(X_test, y_test_cat))

        return history, model

自定义损失函数代码 #

from tensorflow.keras import backend as K

def weighted_categorical_crossentropy(weights):
    """
    A weighted version of keras.objectives.categorical_crossentropy
    By assigning one of the classes a weight of zero, it is possible to ignore that class.
    This allows partial annotation. Typically, class zero is the area that is not annotated,
    and weigh(class=0)=0.
    
    Variables:
        weights: numpy array of shape (C,) where C is the number of classes
    
    Usage example:
        weights = np.array([0.5,2,10]) # Class one at 0.5, class 2 twice the normal weights, class 3 10x.
        loss = weighted_categorical_crossentropy(weights)
        model.compile(loss=loss,optimizer='adam')
        
    Last modifed: 08 April 2024
    """
    
   # weights = K.variable(weights)
    weights = tf.Variable(weights)
        
    def loss(y_true, y_pred):
        # scale predictions so that the class probas of each sample sum to 1
        y_pred /= K.sum(y_pred, axis=-1, keepdims=True)
        # clip to prevent NaN's and Inf's
        y_pred = K.clip(y_pred, K.epsilon(), 1 - K.epsilon())
        # calc
        loss = y_true * K.log(y_pred) * weights
        loss = -K.sum(loss, -1)
        return loss
    
    return loss

优化尝试及疑问 #

尝试以下两个优化步骤后，模型效果有所改善，但出现反常现象：

• 将学习率降至0.0005
• 添加梯度裁剪并使用LossScaleOptimizer

优化后训练IOU持续下降，而验证IOU保持平稳，需解释该现象的原因；同时因单GPU显存仅12GB，batch size调整受限（当前最大为3），寻求进一步的显存优化方向。

内容的提问来源于stack exchange，提问作者user13832811