Per pixel softmax for fully convolutional network

2019-04-06 19:01发布

I'm trying to implement something like a fully convolutional network, where the last convolution layer uses filter size 1x1 and outputs a 'score' tensor. The score tensor has shape [Batch, height, width, num_classes].

My question is, what function in tensorflow can apply softmax operation for each pixel, independent of other pixels. The tf.nn.softmax ops seems not for such purpose.

If there is no such ops available, I guess I have to write one myself.

Thanks!

UPDATE: if I do have to implement myself, I think I may need to reshape the input tensor to [N, num_claees] where N = Batch x width x height, and apply tf.nn.softmax, then reshape it back. Does it make sense?

2条回答
不美不萌又怎样
2楼-- · 2019-04-06 19:21

Reshaping it to 2d and then reshaping it back, like you guessed, is the right approach.

查看更多
够拽才男人
3楼-- · 2019-04-06 19:29

You can use this function.

I found it by searching from GitHub.

import tensorflow as tf

"""
Multi dimensional softmax,
refer to https://github.com/tensorflow/tensorflow/issues/210
compute softmax along the dimension of target
the native softmax only supports batch_size x dimension
"""
def softmax(target, axis, name=None):
    with tf.name_scope(name, 'softmax', values=[target]):
        max_axis = tf.reduce_max(target, axis, keep_dims=True)
        target_exp = tf.exp(target-max_axis)
        normalize = tf.reduce_sum(target_exp, axis, keep_dims=True)
        softmax = target_exp / normalize
        return softmax
查看更多
登录 后发表回答