def resnet_layer(self, x, in_channel, out_channel, stride, dim_match, block_name): conv_name_base = 'res' + block_name + '_branch' bn_name_base = 'bn' + block_name + '_branch' prelu_name_base = 'prelu' + block_name + '_branch' short_cut = x if not dim_match: short_cut = common_layers.conv2d(short_cut, conv_name_base + '1', filter_size=(1, 1), in_channels=in_channel, out_channels=out_channel, strides=stride, bias=False) short_cut = tf.layers.batch_normalization(short_cut, axis=-1, momentum=0.9, training=self.train, name=bn_name_base + '1') x = tf.layers.batch_normalization(x, axis=-1, momentum=0.9, training=self.train, name=bn_name_base + '2a') x = common_layers.conv2d(x, conv_name_base + '2a', (3, 3), in_channel, out_channel, [1, 1], bias=False) x = tf.layers.batch_normalization(x, axis=-1, momentum=0.9, training=self.train, name=bn_name_base + '2b') x = self.prelu_layer(x, name=prelu_name_base + '2b') x = common_layers.conv2d(x, conv_name_base + '2b', (3, 3), out_channel, out_channel, stride, bias=False) res = tf.layers.batch_normalization(x, axis=-1, momentum=0.9, training=self.train, name=bn_name_base + '2c') return tf.add(short_cut, res, name='add_' + block_name)
def se_resnet_layer(self, x, in_channel, out_channel, stride, dim_match, block_name): conv_name_base = 'res_' + block_name + '_branch' bn_name_base = 'bn_' + block_name + '_branch' prelu_name_base = 'prelu_' + block_name + '_branch' se_name_base = 'se_' + block_name + '_branch' short_cut = x if not dim_match: short_cut = common_layers.conv2d(short_cut, conv_name_base + '1', (1, 1), in_channel, out_channel, stride) short_cut = tf.layers.batch_normalization(short_cut, axis=-1, momentum=0.9, training=self.train, name=bn_name_base + '1') x = tf.layers.batch_normalization(x, axis=-1, momentum=0.9, training=self.train, name=bn_name_base + '2a') x = common_layers.conv2d(x, conv_name_base + '2a', (3, 3), in_channel, out_channel, [1, 1]) x = tf.layers.batch_normalization(x, axis=-1, momentum=0.9, training=self.train, name=bn_name_base + '2b') x = self.prelu_layer(x, name=prelu_name_base + '2b') x = common_layers.conv2d(x, conv_name_base + '2b', (3, 3), out_channel, out_channel, stride) x = tf.layers.batch_normalization(x, axis=-1, momentum=0.9, training=self.train, name=bn_name_base + '2c') res = self.se_moudle(x, out_channel, 16, name=se_name_base) return tf.add(short_cut, res, name='add_' + block_name)
def arcface_loss(embedding, labels, out_num, weights=None, s=64., m=0.5, limit_to_pi=True): ''' https://github.com/auroua/InsightFace_TF/blob/master/losses/face_losses.py :param embedding: the input embedding vectors :param labels: the input labels, the shape should be eg: (batch_size, 1) :param s: scalar value default is 64 :param out_num: output class num :param weights: a tf.variable with shape (embedding.shape[-1], out_num) or None to make a new one internally. default = None :param m: the margin value, default is 0.5 :return: the final cacualted output, this output is send into the tf.nn.softmax directly ''' cos_m = math.cos(m) sin_m = math.sin(m) mm = sin_m * m # issue 1 threshold = math.cos(math.pi - m) with tf.variable_scope('arcface_loss'): # inputs and weights norm embedding_norm = tf.norm(embedding, axis=1, keep_dims=True) embedding = tf.div(embedding, embedding_norm, name='norm_embedding') if weights is None: weights = tf.get_variable( name='weights', shape=[embedding.shape[-1].value, out_num], initializer=tf.initializer.glorot_unifrom()) weights_norm = tf.norm(weights, axis=0, keep_dims=True) weights = tf.div(weights, weights_norm, name='norm_weights') # cos(theta+m) cos_t = tf.matmul(embedding, weights, name='cos_t') cos_t2 = tf.square(cos_t, name='cos_2') sin_t2 = tf.subtract(1., cos_t2, name='sin_2') sin_t = tf.sqrt(sin_t2, name='sin_t') cos_mt = s * tf.subtract(tf.multiply(cos_t, cos_m), tf.multiply(sin_t, sin_m), name='cos_mt') if limit_to_pi: # this condition controls the theta+m should in range [0, pi] # 0<=theta+m<=pi # -m<=theta<=pi-m cond_v = cos_t - threshold cond = tf.cast(tf.nn.relu(cond_v, name='if_else'), dtype=tf.bool) keep_val = s * (cos_t - mm) cos_mt_temp = tf.where(cond, cos_mt, keep_val) else: cos_mt_temp = cos_mt mask = tf.one_hot(labels, depth=out_num, name='one_hot_mask') # mask = tf.squeeze(mask, 1) inv_mask = tf.subtract(1., mask, name='inverse_mask') s_cos_t = tf.multiply(s, cos_t, name='scalar_cos_t') output = tf.add(tf.multiply(s_cos_t, inv_mask), tf.multiply(cos_mt_temp, mask), name='arcface_loss_output') return output