keras自定义损失函数即使没有argmax也不会返回梯度

Question

在计算keras的损失之前，我想对y_pred进行一些后处理，因此我决定使用keras自定义损失函数。起初，我使用argmax。我知道argmax没有梯度，所以我将其替换为map_fn，reduce_min和sum等其他函数的组合。但仍然出现错误

  

一个操作具有None用于渐变。请确保所有   您的操作定义了渐变（即可微）。共同   没有渐变的操作：K.argmax，K.round，K.eval。

问题出在哪里？如果问题仍在使用某些特定功能（哪个功能），以及具有梯度的功能等效。使用寄存器Tensorflow梯度，修改keras损失函数或任何其他解决方案？

def my_loss(y_true, y_pred):
  global ind_Batch
  def aaa(x):
    return tf.reduce_min(tf.where(tf.equal(x[0],x[1])))
  def aaa2(xx):
    ww=tf.keras.backend.max(y_pred[ind_Batch,xx,:,:],axis=1)
    return tf.map_fn(aaa,(y_pred[ind_Batch,xx,:,:],ww),dtype=tf.int64)
  def fun_true_false(var,status,mask_):
    delta = tf.SparseTensor(tf.where(mask_), tf.fill((1,tf.count_nonzero(mask_)),status)[0], (384,384)) #assign
    return var+ tf.sparse_tensor_to_dense(delta)
  def my_fun_contour(contour_index):
    ind_F=tf.constant(0)
    def condition(ind_F,num_feature,f):
      return tf.less(ind_F,num_feature+1)
    def body(ind_F,num_feature,f):
      mask=tf.equal(contours,ind_F) #bool mask
      mask_=tf.cast(mask,tf.int32) #int mask
      new_mask=(mask_*argmax_var) #remove unnecessary classes
      #target_class=tf.math.argmax([tf.count_nonzero(tf.equal(new_mask,ind)) for ind in range(1,6)], output_type=tf.dtypes.int32)#find major class
      list_items=[tf.count_nonzero(tf.equal(new_mask,ind)) for ind in range(1,6)]
      target_class=tf.cast(tf.where(tf.equal(list_items,tf.keras.backend.max(list_items)))[0][0],tf.int32)
      for ind in range(6):
        f[ind]=tf.where(tf.equal(ind,target_class+1),fun_true_false(f[ind],1.0,mask_),fun_true_false(f[ind],0.0,mask_))
      return [tf.add(ind_F,1),num_feature,f]
    num_feature=tf.math.reduce_max(tf.math.reduce_max(contours,axis=[1]),axis=0)
    f=[] #a list for appending dimention
    for ind in range(6):
      f.append(np.zeros((384,384),dtype=np.float32))
    res = tf.while_loop(condition,body,[ind_F,num_feature,f])
    mm=tf.transpose(tf.stack(res[-1]),perm=[1,2,0]) #stack f and transpose it to 1,2,0
    return mm
  predict_var=[] #a list for appending batches
  for ind_B in range(Batch_size):
    ind_Batch=ind_B
    argmax_var=tf.cast(tf.map_fn(aaa2,tf.range(384),dtype=tf.int64),tf.int32)
    predict_label=tf.cast(tf.greater(argmax_var,0),tf.int32) 
    contours=tf.contrib.image.connected_components(predict_label)
    predict_var.append(my_fun_contour(ind_B))
  y_pred_=tf.stack(predict_var)
  return bce_jaccard_loss(y_true,y_pred_)

描述我的问题： 我正在进行细分。该代码将找到每个要素（通过tf.contrib.image.connected_components连接的组件），然后计算要素中的主类，最后将主类扩展到整个要素。下图。 谢谢。