I need to convert the following code from tensorflow to pytorch. This code is suppose to generate a position embedding map. Two fully connected layers are used to select features among the prominent ones as well as encode the spatial combinations and feature responses. The feature that is given to function is a convolution feature of size [N, C, H, W], which is output of a VGG16 network at the last conv layer and before the last max pooling layer.
def spatial_aware(input_feature, dimension, trainable, name):
batch, height, width, channel = input_feature.get_shape().as_list()
vec1 = tf.reshape(tf.reduce_mean(input_feature, axis=-1), [-1, height * width])
with tf.variable_scope(name):
weight1 = tf.get_variable(name='weights1', shape=[height * width, int(height * width/2), dimension],
trainable=trainable,
initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.005),
regularizer=tf.contrib.layers.l2_regularizer(0.01))
bias1 = tf.get_variable(name='biases1', shape=[1, int(height * width/2), dimension],
trainable=trainable, initializer=tf.constant_initializer(0.1),
regularizer=tf.contrib.layers.l1_regularizer(0.01))
# vec2 = tf.matmul(vec1, weight1) + bias1
vec2 = tf.einsum('bi, ijd -> bjd', vec1, weight1) + bias1
weight2 = tf.get_variable(name='weights2', shape=[int(height * width / 2), height * width, dimension],
trainable=trainable,
initializer=tf.truncated_normal_initializer(mean=0.0, stddev=0.005),
regularizer=tf.contrib.layers.l2_regularizer(0.01))
bias2 = tf.get_variable(name='biases2', shape=[1, height * width, dimension],
trainable=trainable, initializer=tf.constant_initializer(0.1),
regularizer=tf.contrib.layers.l1_regularizer(0.01))
vec3 = tf.einsum('bjd, jid -> bid', vec2, weight2) + bias2
return vec3