def conv(self, x, kernel_size, num_outputs, name, stride=1, padding="SAME", bias=True, std=1.0):
assert len(x.get_shape()) == 5 # Policies x Batch x Height x Width x Feature
with tf.variable_scope(name):
w = self.create_weight_variable('w', std=std,
shape=(kernel_size, kernel_size, int(x.get_shape()[-1].value), num_outputs))
w = tf.reshape(w, [-1, kernel_size *kernel_size * int(x.get_shape()[-1].value), num_outputs])
x_reshape = tf.reshape(x, (-1, x.get_shape()[2], x.get_shape()[3], x.get_shape()[4]))
patches = tf.extract_image_patches(x_reshape, [1, kernel_size, kernel_size, 1], [1, stride, stride, 1], rates=[1, 1, 1, 1], padding=padding)
final_shape = (tf.shape(x)[0], tf.shape(x)[1], patches.get_shape()[1].value, patches.get_shape()[2].value, num_outputs)
patches = tf.reshape(patches, [tf.shape(x)[0],
-1,
kernel_size * kernel_size * x.get_shape()[-1].value])
if self.indices is None:
ret = tf.matmul(patches, w)
else:
ret = indexed_matmul(patches, w, self.indices)
ret = tf.reshape(ret, final_shape)
self.description += "Convolution layer {} with input shape {} and output shape {}\n".format(name, x.get_shape(), ret.get_shape())
if bias:
b = self.create_bias_variable('b', (1, 1, 1, num_outputs))
if self.indices is not None:
b = tf.gather(b, self.indices)
ret = ret + b
return ret
def dense(self, x, size, name, bias=True, std=1.0):
with tf.variable_scope(name):
w = self.create_weight_variable('w', std=std, shape=(x.get_shape()[-1].value, size))
if self.indices is None:
ret = tf.matmul(x, w)
else:
ret = indexed_matmul(x, w, self.indices)
self.description += "Dense layer {} with input shape {} and output shape {}\n".format(name, x.get_shape(), ret.get_shape())
if bias:
b = self.create_bias_variable('b', (1, size, ))
if self.indices is not None:
b = tf.gather(b, self.indices)
return ret + b
else:
return ret
