какую функцию активации я должен использовать для принудительного округления, подобного поведению

Если вы хотите аппроксимировать округление на реальной линии, вы можете сделать что-то вроде следующего:

def approx_round(x, steepness=1):
    floor_part = tf.floor(x)
    remainder = tf.mod(x, 1)
    return floor_part + tf.sigmoid(steepness*(remainder - 0.5))

На самом деле есть способы зарегистрировать собственные градиенты в Tensorflow (см., например, этот вопрос). Однако я не так знаком с этой частью, так как не так часто использую Keras/TensorFlow.

С точки зрения функции, которая даст вам градиент этого приближения, это будет следующее:

def approx_round_grad(x, steepness=1):
    remainder = tf.mod(x, 1)
    sig = tf.sigmoid(steepness*(remainder - 0.5))
    return sig*(1 - sig)

Чтобы было ясно, это приближение предполагает, что вы используете «достаточно крутой» параметр steepness, поскольку сигмовидная функция не достигает точно 0 или 1, за исключением предела больших аргументов.

Чтобы сделать что-то вроде приближения полусинуса, вы можете использовать следующее:

def approx_round_sin(x, width=0.1):
    if width > 1 or width <= 0:
        raise ValueError('Width must be between zero (exclusive) and one (inclusive)')
    floor_part = tf.floor(x)
    remainder = tf.mod(x, 1)
    return (floor_part + clipped_sin(remainder, width))

def clipped_sin(x, width):
    half_width = width/2
    sin_part = (1 + tf.sin(np.pi*((x-0.5)/width)))/2
    whole = sin_part*tf.cast(tf.abs(x - 0.5) < half_width, tf.float32)
    whole += tf.cast(x > 0.5 + half_width, tf.float32)
    return whole

def approx_round_grad_sin(x, width=0.1):
    if width > 1 or width <= 0:
        raise ValueError('Width must be between zero (exclusive) and one (inclusive)')
    remainder = tf.mod(x, 1)
    return clipped_cos(remainder, width)

def clipped_cos(x, width):
    half_width = width/2
    cos_part = np.pi*tf.cos(np.pi*((x-0.5)/width))/(2*width)
    return cos_part*tf.cast(tf.abs(x - 0.5) < half_width, dtype=tf.float32)

Возможно, вам нужна потеря перекрестной энтропии с функцией softmax tf.nn.softmax_cross_entropy_with_logits_v2, см.

https://www.tensorflow.org/api_docs/python/tf/nn/softmax_cross_entropy_with_logits_v2

Также взгляните на

https://deepnotes.io/softmax-crossentropy