from __future__ import print_function
from __future__ import absolute_import
import keras
import tensorflow as tf
import keras.backend as K
import backbone
weight_decay = 1e-4
class ModelMGPU(keras.Model):
def __init__(self, ser_model, gpus):
pmodel = keras.utils.multi_gpu_model(ser_model, gpus)
self.__dict__.update(pmodel.__dict__)
self._smodel = ser_model
def __getattribute__(self, attrname):
'''Override load and save methods to be used from the serial-model. The
serial-model holds references to the weights in the multi-gpu model.
'''
# return Model.__getattribute__(self, attrname)
if 'load' in attrname or 'save' in attrname:
return getattr(self._smodel, attrname)
return super(ModelMGPU, self).__getattribute__(attrname)
class VladPooling(keras.engine.Layer):
'''
This layer follows the NetVlad, GhostVlad
'''
def __init__(self, mode, k_centers, g_centers=0, **kwargs):
self.k_centers = k_centers
self.g_centers = g_centers
self.mode = mode
super(VladPooling, self).__init__(**kwargs)
def build(self, input_shape):
self.cluster = self.add_weight(shape=[self.k_centers+self.g_centers, input_shape[0][-1]],
name='centers',
initializer='orthogonal')
self.built = True
def compute_output_shape(self, input_shape):
assert input_shape
return (input_shape[0][0], self.k_centers*input_shape[0][-1])
def call(self, x):
# feat : bz x W x H x D, cluster_score: bz X W x H x clusters.
feat, cluster_score = x
num_features = feat.shape[-1]
# softmax normalization to get soft-assignment.
# A : bz x W x H x clusters
max_cluster_score = K.max(cluster_score, -1, keepdims=True)
exp_cluster_score = K.exp(cluster_score - max_cluster_score)
A = exp_cluster_score / K.sum(exp_cluster_score, axis=-1, keepdims = True)
# Now, need to compute the residual, self.cluster: clusters x D
A = K.expand_dims(A, -1) # A : bz x W x H x clusters x 1
feat_broadcast = K.expand_dims(feat, -2) # feat_broadcast : bz x W x H x 1 x D
feat_res = feat_broadcast - self.cluster # feat_res : bz x W x H x clusters x D
weighted_res = tf.multiply(A, feat_res) # weighted_res : bz x W x H x clusters x D
cluster_res = K.sum(weighted_res, [1, 2])
if self.mode == 'gvlad':
cluster_res = cluster_res[:, :self.k_centers, :]
cluster_l2 = K.l2_normalize(cluster_res, -1)
outputs = K.reshape(cluster_l2, [-1, int(self.k_centers) * int(num_features)])
return outputs
def amsoftmax_loss(y_true, y_pred, scale=30, margin=0.35):
y_pred = y_true * (y_pred - margin) + (1 - y_true) * y_pred
y_pred *= scale
return K.categorical_crossentropy(y_true, y_pred, from_logits=True)
def vggvox_resnet2d_icassp(input_dim=(257, 250, 1), num_class=8631, mode='train', args=None):
net=args.net
loss=args.loss
vlad_clusters=args.vlad_cluster
ghost_clusters=args.ghost_cluster
bottleneck_dim=args.bottleneck_dim
aggregation = args.aggregation_mode
mgpu = len(keras.backend.tensorflow_backend._get_available_gpus())
if net == 'resnet34s':
inputs, x = backbone.resnet_2D_v1(input_dim=input_dim, mode=mode)
else:
inputs, x = backbone.resnet_2D_v2(input_dim=input_dim, mode=mode)
# ===============================================
# Fully Connected Block 1
# ===============================================
x_fc = keras.layers.Conv2D(bottleneck_dim, (7, 1),
strides=(1, 1),
activation='relu',
kernel_initializer='orthogonal',
use_bias=True, trainable=True,
kernel_regularizer=keras.regularizers.l2(weight_decay),
bias_regularizer=keras.regularizers.l2(weight_decay),
name='x_fc')(x)
# ===============================================
# Feature Aggregation
# ===============================================
if aggregation == 'avg':
if mode == 'train':
x = keras.layers.AveragePooling2D((1, 5), strides=(1, 1), name='avg_pool')(x)
x = keras.layers.Reshape((-1, bottleneck_dim))(x)
else:
x = keras.layers.GlobalAveragePooling2D(name='avg_pool')(x)
x = keras.layers.Reshape((1, bottleneck_dim))(x)
elif aggregation == 'vlad':
x_k_center = keras.layers.Conv2D(vlad_clusters, (7, 1),
strides=(1, 1),
kernel_initializer='orthogonal',
use_bias=True, trainable=True,
kernel_regularizer=keras.regularizers.l2(weight_decay),
bias_regularizer=keras.regularizers.l2(weight_decay),
name='vlad_center_assignment')(x)
x = VladPooling(k_centers=vlad_clusters, mode='vlad', name='vlad_pool')([x_fc, x_k_center])
elif aggregation == 'gvlad':
x_k_center = keras.layers.Conv2D(vlad_clusters+ghost_clusters, (7, 1),
strides=(1, 1),
kernel_initializer='orthogonal',
use_bias=True, trainable=True,
kernel_regularizer=keras.regularizers.l2(weight_decay),
bias_regularizer=keras.regularizers.l2(weight_decay),
name='gvlad_center_assignment')(x)
x = VladPooling(k_centers=vlad_clusters, g_centers=ghost_clusters, mode='gvlad', name='gvlad_pool')([x_fc, x_k_center])
else:
raise IOError('==> unknown aggregation mode')
# ===============================================
# Fully Connected Block 2
# ===============================================
x = keras.layers.Dense(bottleneck_dim, activation='relu',
kernel_initializer='orthogonal',
use_bias=True, trainable=True,
kernel_regularizer=keras.regularizers.l2(weight_decay),
bias_regularizer=keras.regularizers.l2(weight_decay),
name='fc6')(x)
# ===============================================
# Softmax Vs AMSoftmax
# ===============================================
if loss == 'softmax':
y = keras.layers.Dense(num_class, activation='softmax',
kernel_initializer='orthogonal',
use_bias=False, trainable=True,
kernel_regularizer=keras.regularizers.l2(weight_decay),
bias_regularizer=keras.regularizers.l2(weight_decay),
name='prediction')(x)
trnloss = 'categorical_crossentropy'
elif loss == 'amsoftmax':
x_l2 = keras.layers.Lambda(lambda x: K.l2_normalize(x, 1))(x)
y = keras.layers.Dense(num_class,
kernel_initializer='orthogonal',
use_bias=False, trainable=True,
kernel_constraint=keras.constraints.unit_norm(),
kernel_regularizer=keras.regularizers.l2(weight_decay),
bias_regularizer=keras.regularizers.l2(weight_decay),
name='prediction')(x_l2)
trnloss = amsoftmax_loss
else:
raise IOError('==> unknown loss.')
if mode == 'eval':
y = keras.layers.Lambda(lambda x: keras.backend.l2_normalize(x, 1))(x)
model = keras.models.Model(inputs, y, name='vggvox_resnet2D_{}_{}'.format(loss, aggregation))
if mode == 'train':
if mgpu > 1:
model = ModelMGPU(model, gpus=mgpu)
# set up optimizer.
if args.optimizer == 'adam': opt = keras.optimizers.Adam(lr=1e-3)
elif args.optimizer =='sgd': opt = keras.optimizers.SGD(lr=0.1, momentum=0.9, decay=0.0, nesterov=True)
else: raise IOError('==> unknown optimizer type')
model.compile(optimizer=opt, loss=trnloss, metrics=['acc'])
return model
