基于TensorFlow/Keras的像素级角度距离损失函数优化实现问询

我尝试基于一篇实用论文，在TensorFlow/Keras中实现用于方向图预测的损失函数。作者通过预测每个像素（输出的每个通道）的正弦和余弦值，定义了如下距离度量公式：

θ^(1+δ) = (arccos(cosα · cosβ + sinα · sinβ))^(1+δ)
其中δ=0.2，对应的梯度为(1 + δ) · θ^δ。

已知α对应真实标签张量y_true，β对应预测张量y_pred，二者形状均为(batch, height, width, channels)。我目前用嵌套循环实现了该损失函数，但完全没有做优化，而且因为缺乏ML经验，不确定Keras能否对这个实现进行反向传播。

我想了解有没有比以下代码更优的实现方式，代码里的min和max函数是用来将值裁剪到区间[1e-6, 1-1e-6]的：

def angle_distance_loss(y_true,y_pred):
    """
    Lproposed = (arccos(cosα · cosβ + sinα · sinβ))^(1+δ)
    """

    batch, height, width, channels = y_true.shape
    cos_c = 0
    sin_c = 1
    l = 0  # 补充原代码遗漏的变量初始化
    for batch_i in range(batch):
        for h_j in range(height):
            for w_k in range(width):
                yt_cos = y_true[batch_i][h_j][w_k][cos_c]
                yt_sin = y_true[batch_i][h_j][w_k][sin_c]

                yp_cos = y_pred[batch_i][h_j][w_k][cos_c]
                yp_sin = y_pred[batch_i][h_j][w_k][sin_c]

                dot_product = yt_cos * yp_cos + yt_sin * yp_sin
                clipped_dot = max(10**-6, min(dot_product, 1-10**-6))
                l += math.acos(clipped_dot) ** 1.2
                
    return l / (batch * width * height)  # 修正原代码除法优先级，保证平均逻辑正确

希望能得到关于优化实现、保证反向传播正常工作的建议 :)

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