Windows本地安装sonnet失败，所以只试验了Gaussian Attention，Spatial Transformer可用sonnet实现，只是sonnet安装失败。
import tensorflowas tf
#import sonnet as snt
import numpyas np
import matplotlib.pyplotas plt
def gaussian_mask(u, s, d, R, C):
"""
:param u: tf.Tensor, centre of the first Gaussian.
:param s: tf.Tensor, standard deviation of Gaussians.
:param d: tf.Tensor, shift between Gaussian centres.
:param R: int, number of rows in the mask, there is one Gaussian per row.
:param C: int, number of columns in the mask.
"""
# indices to create centres
R= tf.to_float(tf.reshape(tf.range(R), (1, 1, R)))
C= tf.to_float(tf.reshape(tf.range(C), (1, C, 1)))
centres= u[np.newaxis, :, np.newaxis]+ R * d
column_centres= C - centres
mask= tf.exp(-.5 * tf.square(column_centres/ s))
# we add eps for numerical stability
normalised_mask= mask/ (tf.reduce_sum(mask, 1, keep_dims=True)+ 1e-8)
return normalised_mask
def gaussian_glimpse(img_tensor, transform_params, crop_size):
"""
:param img_tensor: tf.Tensor of size (batch_size, Height, Width, channels)
:param transform_params: tf.Tensor of size (batch_size, 6), where params are (mean_y, std_y, d_y, mean_x, std_x, d_x) specified in pixels.
:param crop_size): tuple of 2 ints, size of the resulting crop
"""
# parse arguments
h, w= crop_size
H, W= img_tensor.shape.as_list()[1:3]
split_ax= transform_params.shape.ndims-1
uy, sy, dy, ux, sx, dx= tf.split(transform_params, 6, split_ax)
# create Gaussian masks, one for each axis
Ay= gaussian_mask(uy, sy, dy, h, H)
Ax= gaussian_mask(ux, sx, dx, w, W)
# extract glimpse
glimpse= tf.matmul(tf.matmul(Ay, img_tensor, adjoint_a=True), Ax)
return glimpse
img_size= 10, 10
glimpse_size= 5, 5
# Create a random image with a square
x= abs(np.random.randn(1, *img_size))* .3
x[0, 3:6, 3:6]= 1
crop= x[0, 1:8, 1:8]# contains the square
tf.reset_default_graph()
# placeholders
tx= tf.placeholder(tf.float32, x.shape, 'image')
tu= tf.placeholder(tf.float32, [1], 'u')
ts= tf.placeholder(tf.float32, [1], 's')
td= tf.placeholder(tf.float32, [1], 'd')
stn_params= tf.placeholder(tf.float32, [1, 4], 'stn_params')
# Gaussian Attention
gaussian_att_params= tf.concat([tu, ts, td, tu, ts, td], -1)
gaussian_glimpse_expr= gaussian_glimpse(tx, gaussian_att_params, glimpse_size)
# Spatial Transformer
#stn_glimpse_expr = spatial_transformer(tx, stn_params, glimpse_size)
sess= tf.Session()
# extract a Gaussian glimpse
u= 2
s= .5
d= 1
u, s, d= (np.asarray([i])for iin (u, s, d))
gaussian_crop= sess.run(gaussian_glimpse_expr, feed_dict={tx: x, tu: u, ts: s, td: d})
# extract STN glimpse
# transform = [.4, -.1, .4, -.1]
# transform = np.asarray(transform).reshape((1, 4))
# stn_crop = sess.run(stn_glimpse_expr, {tx: x, stn_params: transform})
# plots
fig, axes= plt.subplots(1, 3, figsize=(12, 3))
titles= ['Input Image', 'Crop', 'Gaussian Att']#, 'STN']
imgs= [x, crop, gaussian_crop]#, stn_crop]
for ax, title, imgin zip(axes, titles, imgs):
ax.imshow(img.squeeze(), cmap='gray', vmin=0., vmax=1.)
ax.set_title(title)
ax.xaxis.set_visible(False)
ax.yaxis.set_visible(False)
plt.show()
