(inputs, name='attention', reuse=False)
| 64 | |
| 65 | |
| 66 | def self_attention(inputs, name='attention', reuse=False): |
| 67 | with tf.variable_scope(name, reuse=reuse): |
| 68 | h, w = tf.shape(inputs)[1], tf.shape(inputs)[2] |
| 69 | bs, _, _, ch = inputs.get_shape().as_list() |
| 70 | f = slim.convolution2d(inputs, ch//8, [1, 1], activation_fn=None) |
| 71 | g = slim.convolution2d(inputs, ch//8, [1, 1], activation_fn=None) |
| 72 | s = slim.convolution2d(inputs, 1, [1, 1], activation_fn=None) |
| 73 | f_flatten = tf.reshape(f, shape=[f.shape[0], -1, f.shape[-1]]) |
| 74 | g_flatten = tf.reshape(g, shape=[g.shape[0], -1, g.shape[-1]]) |
| 75 | beta = tf.matmul(f_flatten, g_flatten, transpose_b=True) |
| 76 | beta = tf.nn.softmax(beta) |
| 77 | |
| 78 | s_flatten = tf.reshape(s, shape=[s.shape[0], -1, s.shape[-1]]) |
| 79 | att_map = tf.matmul(beta, s_flatten) |
| 80 | att_map = tf.reshape(att_map, shape=[bs, h, w, 1]) |
| 81 | gamma = tf.get_variable("gamma", [1], initializer=tf.constant_initializer(0.0)) |
| 82 | output = att_map * gamma + inputs |
| 83 | |
| 84 | return att_map, output |
| 85 | |
| 86 | |
| 87 |
nothing calls this directly
no outgoing calls
no test coverage detected