Softpix por píxel para una red totalmente convolucional

Estoy tratando de implementar algo así como una red totalmente convolucional, donde la última capa de convolución usa un tamaño de filtro 1x1 y genera un tensor de 'puntuación'. El tensor de puntuación tiene forma [Lote, altura, ancho, num_classes].

Mi pregunta es, qué función en tensorflow puede aplicar la operación softmax para cada píxel, independientemente de otros píxeles. Las operaciones tf.nn.softmax no parecen ser para tal propósito.

Si no hay tales operaciones disponibles, supongo que tengo que escribir una yo mismo.

¡Gracias!

ACTUALIZACIÓN: si tengo que implementarme, creo que es posible que deba cambiar la forma del tensor de entrada a [N, num_claees] donde N = Lote x ancho x altura, y aplicar tf.nn.softmax, luego volver a darle forma. ¿Tiene sentido?

Respuestas a la pregunta(2)

Su respuesta a la pregunta