network.py

class ResNet_101(object):
	def __init__(self, inputs, num_classes, phase):
		self.inputs = inputs
		self.num_classes = num_classes
		self.channel_axis = 3
		self.phase = phase # train (True) or test (False), for BN layers in the decoder
		self.build_network()

	def build_network(self):
		self.encoding = self.build_encoder()
		outputs = self._avg_pool(self.encoding, pool_size=[1,1], stride=1, name='poo5')
		outputs = self._conv2d(outputs, 1, 1000, 1, name='fc6', biased=True)

	def build_encoder(self):
		scope_name = 'resnet_v1_101'
		with tf.variable_scope(scope_name) as scope:
			outputs = self._start_block('conv1')
			print("after start block:", outputs.shape)
			with tf.variable_scope('block1') as scope:
				outputs = self._bottleneck_resblock(outputs, 256, 'unit_1',	identity_connection=False)
				outputs = self._bottleneck_resblock(outputs, 256, 'unit_2')
				outputs = self._bottleneck_resblock(outputs, 256, 'unit_3')
				print("after block1:", outputs.shape)
			with tf.variable_scope('block2') as scope:
				outputs = self._bottleneck_resblock(outputs, 512, 'unit_1',	half_size=True, identity_connection=False)
				for i in six.moves.range(2, 5):
					outputs = self._bottleneck_resblock(outputs, 512, 'unit_%d' % i)
				print("after block2:", outputs.shape)
			with tf.variable_scope('block3') as scope:
				outputs = self._bottleneck_resblock(outputs, 1024, 'unit_1',half_size=True,	identity_connection=False)
				num_layers_block3 = 23
				for i in six.moves.range(2, num_layers_block3+1):
					outputs = self._bottleneck_resblock(outputs, 1024, 'unit_%d' % i)
				print("after block3:", outputs.shape)
			with tf.variable_scope('block4') as scope:
				outputs = self._bottleneck_resblock(outputs, 2048, 'unit_1',half_size=True, identity_connection=False)
				outputs = self._bottleneck_resblock(outputs, 2048, 'unit_2')
				outputs = self._bottleneck_resblock(outputs, 2048, 'unit_3')
				print("after block4:", outputs.shape)

				return outputs

	# blocks
	def _start_block(self, name):
		outputs = self._conv2d(self.inputs, 7, 64, 2, name=name)
		outputs = self._batch_norm(outputs, name=name, is_training=False, activation_fn=tf.nn.relu)
		outputs = self._max_pool2d(outputs, 3, 2, name='pool1')
		return outputs

	def _bottleneck_resblock(self, x, num_o, name, half_size=False, identity_connection=True):
		first_s = 2 if half_size else 1
		assert num_o % 4 == 0, 'Bottleneck number of output ERROR!'
		# branch1
		if not identity_connection:
			o_b1 = self._conv2d(x, 1, num_o, first_s, name='%s/bottleneck_v1/shortcut' % name)
			o_b1 = self._batch_norm(o_b1, name='%s/bottleneck_v1/shortcut' % name, is_training=False, activation_fn=None)
		else:
			o_b1 = x
		# branch2
		o_b2a = self._conv2d(x, 1, num_o / 4, first_s, name='%s/bottleneck_v1/conv1' % name)
		o_b2a = self._batch_norm(o_b2a, name='%s/bottleneck_v1/conv1' % name, is_training=False, activation_fn=tf.nn.relu)

		o_b2b = self._conv2d(o_b2a, 3, num_o / 4, 1, name='%s/bottleneck_v1/conv2' % name)
		o_b2b = self._batch_norm(o_b2b, name='%s/bottleneck_v1/conv2' % name, is_training=False, activation_fn=tf.nn.relu)

		o_b2c = self._conv2d(o_b2b, 1, num_o, 1, name='%s/bottleneck_v1/conv3' % name)
		o_b2c = self._batch_norm(o_b2c, name='%s/bottleneck_v1/conv3' % name, is_training=False, activation_fn=None)
		# add
		outputs = self._add([o_b1,o_b2c], name='%s/bottleneck_v1/add' % name)
		# relu
		outputs = self._relu(outputs, name='%s/bottleneck_v1/relu' % name)
		return outputs

	# layers
	def _conv2d(self, x, kernel_size, num_o, stride, name, biased=False):
		"""
		Conv2d without BN or relu.
		"""
		num_x = x.shape[self.channel_axis].value
		with tf.variable_scope(name) as scope:
			w = tf.get_variable('weights', shape=[kernel_size, kernel_size, num_x, num_o])
			s = [1, stride, stride, 1]
			o = tf.nn.conv2d(x, w, s, padding='SAME')
			if biased:
				b = tf.get_variable('biases', shape=[num_o])
				o = tf.nn.bias_add(o, b)
			return o

	def _relu(self, x, name):
		return tf.nn.relu(x, name=name)

	def _add(self, x_l, name):
		return tf.add_n(x_l, name=name)

	def _max_pool2d(self, x, kernel_size, stride, name):
		k = [1, kernel_size, kernel_size, 1]
		s = [1, stride, stride, 1]
		return tf.nn.max_pool(x, k, s, padding='SAME', name=name)

	def _batch_norm(self, x, name, is_training, activation_fn, trainable=False):
		# For a small batch size, it is better to keep
		# the statistics of the BN layers (running means and variances) frozen,
		# and to not update the values provided by the pre-trained model by setting is_training=False.
		# Note that is_training=False still updates BN parameters gamma (scale) and beta (offset)
		# if they are presented in var_list of the optimiser definition.
		# Set trainable = False to remove them from trainable_variables.
		with tf.variable_scope(name+'/BatchNorm') as scope:
			o = tf.contrib.layers.batch_norm(
				x,
				scale=True,
				activation_fn=activation_fn,
				is_training=is_training,
				trainable=trainable,
				scope=scope)
			return o
	def _avg_pool(self, x, pool_size=[2,2], stride=2, name='pool' ,padding='VALID'):
		return tf.layers.average_pooling2d(inputs=x, pool_size=pool_size, strides=stride, name=name,padding=padding,)


class ResNet_101_Extra(object):
	def __init__(self, inputs, num_classes, phase):
		self.inputs = inputs
		self.num_classes = num_classes
		self.channel_axis = 3
		self.phase = phase # train (True) or test (False), for BN layers in the decoder
		self.build_network()

	def build_network(self):
		self.encoding = self.build_encoder()
		outputs = self._avg_pool(self.encoding, pool_size=[1,1], stride=1, name='poo5_extra')
		outputs = self._conv2d(outputs, 1, 1000, 1, name='fc6_extra', biased=True)

	def build_encoder(self):
		scope_name = 'resnet_v1_101_extra'
		with tf.variable_scope(scope_name) as scope:
			outputs = self._start_block('conv1')
			print("after start block:", outputs.shape)
			with tf.variable_scope('block1') as scope:
				outputs = self._bottleneck_resblock(outputs, 256, 'unit_1',	identity_connection=False)
				outputs = self._bottleneck_resblock(outputs, 256, 'unit_2')
				outputs = self._bottleneck_resblock(outputs, 256, 'unit_3')
				print("after block1:", outputs.shape)
			with tf.variable_scope('block2') as scope:
				outputs = self._bottleneck_resblock(outputs, 512, 'unit_1',	half_size=True, identity_connection=False)
				for i in six.moves.range(2, 5):
					outputs = self._bottleneck_resblock(outputs, 512, 'unit_%d' % i)
				print("after block2:", outputs.shape)
			with tf.variable_scope('block3') as scope:
				outputs = self._bottleneck_resblock(outputs, 1024, 'unit_1',half_size=True,	identity_connection=False)
				num_layers_block3 = 23
				for i in six.moves.range(2, num_layers_block3+1):
					outputs = self._bottleneck_resblock(outputs, 1024, 'unit_%d' % i)
				print("after block3:", outputs.shape)
			with tf.variable_scope('block4') as scope:
				outputs = self._bottleneck_resblock(outputs, 2048, 'unit_1',half_size=True, identity_connection=False)
				outputs = self._bottleneck_resblock(outputs, 2048, 'unit_2')
				outputs = self._bottleneck_resblock(outputs, 2048, 'unit_3')
				print("after block4:", outputs.shape)
			with tf.variable_scope('block5') as scope:
				outputs = self._bottleneck_resblock(outputs, 2048, 'unit_1', half_size=True,
													identity_connection=False)
				outputs = self._bottleneck_resblock(outputs, 2048, 'unit_2')
				outputs = self._bottleneck_resblock(outputs, 2048, 'unit_3')
				print("after block5:", outputs.shape)
			with tf.variable_scope('block6') as scope:
				outputs = self._bottleneck_resblock(outputs, 2048, 'unit_1', half_size=True,
													identity_connection=False)
				outputs = self._bottleneck_resblock(outputs, 2048, 'unit_2')
				outputs = self._bottleneck_resblock(outputs, 2048, 'unit_3')
				print("after block5:", outputs.shape)
			with tf.variable_scope('block7') as scope:
				outputs = self._bottleneck_resblock(outputs, 2048, 'unit_1', half_size=True,
													identity_connection=False)
				outputs = self._bottleneck_resblock(outputs, 2048, 'unit_2')
				outputs = self._bottleneck_resblock(outputs, 2048, 'unit_3')
				print("after block7:", outputs.shape)

				return outputs