selcnn.py

import tensorflow as tf
import numpy as np

from utils import variable_on_cpu, variable_with_weight_decay

def non2zero(l):
  return [0 if i==None else i for i in l] 

class SelCNN:
	def __init__(self, scope, gt_M_sz):
		"""
		selCNN network class. Initialize graph.

		Args: 
			name: string, name of the network.
			vgg_conv_layer: tensor, eithedoner conv4_3 or cnv5_3 layer
				of a pretrained vgg16 network
		"""
		# Initialize network
		self.scope = scope
	
		# COnstract a packed 1 x 28 x 28 x 512 dim input
		self._get_input() 
		

		self.variables = []
		# Best params for lselCNN: lr = 1e-6, decay_step = 500, decay rate 0.5
		# with premutated img. rms loss. 
		self.params = {
		'dropout_rate': 0.3,
		'k_size': [3, 3, 512, 1],
		'wd': 0.0,
		'lr_initial': 1e-6, # 1e-8 gives 438 after 200 steps, 1e-7 gives better maps?
		'lr_decay_steps': 500 ,
		'lr_decay_rate': 0.5
		}
		with tf.name_scope(scope) as scope:
			self.pre_M = self._get_pre_M()
			self.gt_M = tf.placeholder(tf.float32, shape=gt_M_sz)
		self.pre_M_size = self.pre_M.get_shape().as_list()[1:3]
		
	def _get_input(self):
		self.input_maps = tf.placeholder(shape=[1,224,224,512], dtype=tf.float32)
		self.feature_maps = [self.input_maps[...,i] for i in range(512)]
		

	def _get_pre_M(self):
		"""Build the sel-CNN graph and returns predicted Heat map."""
		input_maps = tf.pack(self.feature_maps, axis=-1)
		dropout_layer = tf.nn.dropout(input_maps, self.params['dropout_rate'])

		# Conv layer with bias 
		kernel = variable_with_weight_decay(self.scope, 'kernel',\
							self.params['k_size'], wd = self.params['wd'])
		conv = tf.nn.conv2d(dropout_layer, kernel, [1,1,1,1], 'SAME')
		bias = variable_on_cpu(self.scope,'biases', [1], tf.constant_initializer(0.1))
		pre_M = tf.nn.bias_add(conv, bias)
		self.variables += [kernel, bias]
		# Subtract mean 
		pre_M -= tf.reduce_mean(pre_M)
		#pre_M# /= tf.reduce_max(pre_M)
		return pre_M


	def train_op(self, global_step, add_regulizer=False):
		""" Train the network on the fist frame. 

		Args:
			gt_M_sz: tuple, shape identical to self.pre_M,
				Ground truth heatmap.
			add_regulizer: bool, True for adding L2 regulizer of the 
				kernel variables of the conv layer.

		Returns:
			train_op:
			total_losses:
			lr:
		"""
		pre_shape = self.pre_M.get_shape().as_list()[1:]
		gt_shape = self.gt_M.get_shape().as_list()[1:]
		assert gt_shape == pre_shape, 'Shapes are not compatiable! gt_M : {0}, pre_M : {1}'.format(gt_shape, pre_shape)
		
		# Root mean square loss
		rms_loss = tf.reduce_mean(tf.squared_difference(self.gt_M, self.pre_M))
		
		tf.add_to_collection('losses', rms_loss)
		# Use vanila SGD with exponentially decayed learning rate
		# Decayed_learning_rate = learning_rate *
		#                decay_rate ^ (global_step / decay_steps)
		lr = tf.train.exponential_decay(
			self.params['lr_initial'], 
			global_step, 
			self.params['lr_decay_steps'], 
			self.params['lr_decay_rate'] , 
			name='lr')

		# Vanilia SGD with dexp decay learning rate
		optimizer = tf.train.GradientDescentOptimizer(lr)

		if add_regulizer:
			# Add L2 regularzer losses
			total_losses = tf.add_n(tf.get_collection(tf.GraphKeys.LOSSES), 'total_losses')
		else:
			total_losses = rms_loss
		train_op = optimizer.minimize(total_losses, var_list = self.variables, global_step=global_step)
		self.loss = total_losses
		return train_op, total_losses, lr, optimizer

	def sel_feature_maps(self, sess, feed_dict, num_sel):
		""" 
		Selects saliency feature maps. 
		The change of the Loss function by the permutation
		of the feature maps dF, can be computed by a 
		two-order Taylor expansions.

		Further simplication can be made by only compute
		the diagonol part of the Hessian matrix.

		S = - partial(L)/patial(F) * F 
			+ 0.5 * sencondOrderPartial(L)/F

		Args:
			sel_maps: tensor, conv layer of vgg.
			num_sel: int, number of selected maps.

		Returns:
			idx: list, indexes of selected maps
		"""
		# Compute first derevatives w.r.t each feature maps
		grads = non2zero(tf.gradients(self.loss, self.feature_maps))

		# Compute diagnol part of Hessian which are the second derevatives
		# of Loss_x w.r.t x
		H_diag = [non2zero(tf.gradients(grads[i], self.feature_maps[i]))[0] for i in range(512)]

		# Compute the significance vector, with each element stand for 
		# the score of each feature map
		S = [tf.reduce_sum(-tf.mul(grads[i], self.feature_maps[i])) \
			+ 0.5 * tf.reduce_sum(tf.mul(H_diag[i], self.feature_maps[i]**2)) for i in range(512)]
		S_tensor = tf.pack(S, axis=0) # shape (512,)

		signif_v = sess.run([S_tensor], feed_dict=feed_dict)

		# Retrieve the top-num_sel feature maps and corresponding idx
		idxs = sorted(range(len(signif_v)), key=lambda i: signif_v[i])[-num_sel:]
		#best_maps = vgg_maps[...,idxs]
		#print('Selected maps shape: {0}'.format(best_maps.shape)) # e.g.(1, 28, 28, 384)
		return idxs