None_local_block in TensorFlow

this block are not offcially implement of non_local_block in paper

Non-local Neural Networks https://arxiv.org/pdf/1711.07971.pdf

this code just suit for 2D images data, you can change this code easilly to suit for 3D or 4D data

Whole code you can visit https://github.com/jack-Dong/None_local_block

code

def non_local_block(input_tensor, computation_compression=2, mode='dot'): if mode not in ['gaussian', 'embedded', 'dot', 'concatenate']: raise ValueError('`mode` must be one of `gaussian`, `embedded`, `dot`,`concatenate`') input_shape = input_tensor.get_shape().as_list() print(input_shape) batchsize, dim1, dim2, channels = input_shape if mode == 'gaussian': # Gaussian instantiation x1 = tf.reshape(input_tensor, shape=[-1, dim1 * dim2, channels]) x2 = tf.reshape(input_tensor, shape=[-1, dim1 * dim2, channels]) f = tf.matmul(x1, x2, transpose_b=True) f = tf.reshape(f, shape=[-1, dim1 * dim2 * dim1 * dim2]) f = tf.nn.softmax(f, axis=-1) f = tf.reshape(f, shape=[-1, dim1 * dim2, dim1 * dim2]) print("gaussian=", f) elif mode == 'dot': theta = conv2d(input_tensor, channels, channels // 2, 1) # add BN relu layer before conv will speed up training theta = tf.reshape(theta, shape=[-1, dim1 * dim2, channels // 2]) phi = conv2d(input_tensor, channels, channels // 2, 1) phi = tf.reshape(phi, shape=[-1, dim1 * dim2, channels // 2]) f = tf.matmul(theta, phi, transpose_b=True) # scale the values to make it size invarian t f = f / (dim1 * dim2 * channels) print("dot f=", f) elif mode == 'concatenate': # this operation cost too much memory, so make sure you input a small resolution feature map like(14X14 7X7) theta = conv2d(input_tensor, channels, channels // 2, 1) theta = tf.reshape(theta, shape=[-1, dim1 * dim2, channels // 2]) phi = conv2d(input_tensor, channels, channels // 2, 1) phi = tf.reshape(phi, shape=[-1, dim1 * dim2, channels // 2]) theta_splits = tf.split(theta, dim1 * dim2, 1) phi_splits = tf.split(phi, dim1 * dim2, 1) theta_split_shape = tf.shape(theta[0]) print("theta_split_shape", theta_split_shape) initial = tf.constant(1.0 / channels, shape=[channels, 1]) print('initial', initial) W_concat = tf.Variable(initial) print("W_concat", W_concat) f_matrix = [] for i in range(dim1 * dim2): for j in range(dim1 * dim2): print(i, ' ', j) tmp = tf.concat([theta_splits[i], phi_splits[j]], 2) tmp = tf.reshape(tmp, shape=[-1, channels]) # print(tmp) tmp = tf.matmul(tmp, W_concat) print(tmp) f_matrix.append(tmp) f_matrix_tensor = tf.stack(f_matrix, axis=2) print('f_matrix_tensor', f_matrix_tensor) f = tf.reshape(f_matrix_tensor, shape=[-1, dim1 * dim2, dim1 * dim2]) f = f / (dim1 * dim2 * channels) print("concatenate f=", f) else: # Embedded Gaussian instantiation theta = conv2d(input_tensor, channels, channels // 2, 1) theta = tf.reshape(theta, shape=[-1, dim1 * dim2, channels // 2]) phi = conv2d(input_tensor, channels, channels // 2, 1) phi = tf.reshape(phi, shape=[-1, dim1 * dim2, channels // 2])<p> if computation_compression > 1: phi = tf.layers.max_pooling1d(phi, pool_size=2, strides=computation_compression, padding='SAME') print('phi', phi) f = tf.matmul(theta, phi, transpose_b=True) phi_shape = phi.get_shape().as_list() f = tf.reshape(f, shape=[-1, dim1 * dim2 * phi_shape[1]]) f = tf.nn.def non_local_block(input_tensor, computation_compression=2, mode='dot'): if mode not in ['gaussian', 'embedded', 'dot', 'concatenate']: raise ValueError('`mode` must be one of `gaussian`, `embedded`, `dot`,`concatenate`') input_shape = input_tensor.get_shape().as_list() print(input_shape) batchsize, dim1, dim2, channels = input_shape if mode == 'gaussian': # Gaussian instantiation x1 = tf.reshape(input_tensor, shape=[-1, dim1 * dim2, channels]) x2 = tf.reshape(input_tensor, shape=[-1, dim1 * dim2, channels]) f = tf.matmul(x1, x2, transpose_b=True) f = tf.reshape(f, shape=[-1, dim1 * dim2 * dim1 * dim2]) f = tf.nn.softmax(f, axis=-1) f = tf.reshape(f, shape=[-1, dim1 * dim2, dim1 * dim2]) print("gaussian=", f) elif mode == 'dot': theta = conv2d(input_tensor, channels, channels // 2, 1) # add BN relu layer before conv will speed up training theta = tf.reshape(theta, shape=[-1, dim1 * dim2, channels // 2]) phi = conv2d(input_tensor, channels, channels // 2, 1) phi = tf.reshape(phi, shape=[-1, dim1 * dim2, channels // 2]) f = tf.matmul(theta, phi, transpose_b=True) # scale the values to make it size invarian t f = f / (dim1 * dim2 * channels) print("dot f=", f) elif mode == 'concatenate': # this operation cost too much memory, so make sure you input a small resolution feature map like(14X14 7X7) theta = conv2d(input_tensor, channels, channels // 2, 1) theta = tf.reshape(theta, shape=[-1, dim1 * dim2, channels // 2]) phi = conv2d(input_tensor, channels, channels // 2, 1) phi = tf.reshape(phi, shape=[-1, dim1 * dim2, channels // 2]) <p> theta_splits = tf.split(theta, dim1 * dim2, 1) phi_splits = tf.split(phi, dim1 * dim2, 1) theta_split_shape = tf.shape(theta[0]) print("theta_split_shape", theta_split_shape) initial = tf.constant(1.0 / channels, shape=[channels, 1]) print('initial', initial) W_concat = tf.Variable(initial) print("W_concat", W_concat) f_matrix = [] for i in range(dim1 * dim2): for j in range(dim1 * dim2): print(i, ' ', j) tmp = tf.concat([theta_splits[i], phi_splits[j]], 2) tmp = tf.reshape(tmp, shape=[-1, channels]) # print(tmp) tmp = tf.matmul(tmp, W_concat) print(tmp) f_matrix.append(tmp) f_matrix_tensor = tf.stack(f_matrix, axis=2) print('f_matrix_tensor', f_matrix_tensor) f = tf.reshape(f_matrix_tensor, shape=[-1, dim1 * dim2, dim1 * dim2]) f = f / (dim1 * dim2 * channels) print("concatenate f=", f)<p> else: # Embedded Gaussian instantiation theta = conv2d(input_tensor, channels, channels // 2, 1) theta = tf.reshape(theta, shape=[-1, dim1 * dim2, channels // 2]) phi = conv2d(input_tensor, channels, channels // 2, 1) phi = tf.reshape(phi, shape=[-1, dim1 * dim2, channels // 2]) if computation_compression > 1: phi = tf.layers.max_pooling1d(phi, pool_size=2, strides=computation_compression, padding='SAME') print('phi', phi) f = tf.matmul(theta, phi, transpose_b=True) phi_shape = phi.get_shape().as_list() f = tf.reshape(f, shape=[-1, dim1 * dim2 * phi_shape[1]]) f = tf.nn.softmax(f, axis=-1) f = tf.reshape(f, shape=[-1, dim1 * dim2, phi_shape[1]]) print("Embedded f=", f) g = conv2d(input_tensor, channels, channels // 2, 1) g = tf.reshape(g, shape=[-1, dim1 * dim2, channels // 2]) if computation_compression > 1 and mode == 'embedded': g = tf.layers.max_pooling1d(g, pool_size=2, strides=computation_compression, padding='SAME') print('g', g) y = tf.matmul(f, g) print('y=', y) y = tf.reshape(y, shape=[-1, dim1, dim2, channels // 2]) y = conv2d(y, channels // 2, channels, kernel_size=3) print('y=', y) residual = input_tensor + y return residualsoftmax(f, axis=-1) f = tf.reshape(f, shape=[-1, dim1 * dim2, phi_shape[1]]) print("Embedded f=", f) g = conv2d(input_tensor, channels, channels // 2, 1) g = tf.reshape(g, shape=[-1, dim1 * dim2, channels // 2]) if computation_compression > 1 and mode == 'embedded': g = tf.layers.max_pooling1d(g, pool_size=2, strides=computation_compression, padding='SAME') print('g', g) y = tf.matmul(f, g) print('y=', y) y = tf.reshape(y, shape=[-1, dim1, dim2, channels // 2]) y = conv2d(y, channels // 2, channels, kernel_size=3) print('y=', y) residual = input_tensor + y return residual