MCPcopy Create free account
hub / github.com/SystemErrorWang/White-box-Cartoonization / self_attention

Function self_attention

train_code/layers.py:66–84  ·  view source on GitHub ↗
(inputs, name='attention', reuse=False)

Source from the content-addressed store, hash-verified

64
65
66def 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

Callers

nothing calls this directly

Calls

no outgoing calls

Tested by

no test coverage detected