def __init__(self):
"""Passes convolutional features through embedding network.
"""
super(ConvEmbedder, self).__init__()
self.num_steps = CONFIG.DATA.NUM_STEPS
conv_params = CONFIG.MODEL.CONV_EMBEDDER_MODEL.CONV_LAYERS
fc_params = CONFIG.MODEL.CONV_EMBEDDER_MODEL.FC_LAYERS
use_bn = CONFIG.MODEL.CONV_EMBEDDER_MODEL.USE_BN
l2_reg_weight = CONFIG.MODEL.L2_REG_WEIGHT
embedding_size = CONFIG.MODEL.CONV_EMBEDDER_MODEL.EMBEDDING_SIZE
cap_scalar = CONFIG.MODEL.CONV_EMBEDDER_MODEL.CAPACITY_SCALAR
conv_params = [(cap_scalar * x[0], x[1], x[2]) for x in conv_params]
fc_params = [(cap_scalar * x[0], x[1]) for x in fc_params]
conv_bn_activations = get_conv_bn_layers(conv_params,
use_bn,
conv_dims=3)
self.conv_layers = conv_bn_activations[0]
self.bn_layers = conv_bn_activations[1]
self.activations = conv_bn_activations[2]
self.fc_layers = get_fc_layers(fc_params)
self.embedding_layer = layers.Dense(
embedding_size,
kernel_regularizer=regularizers.l2(l2_reg_weight),
bias_regularizer=regularizers.l2(l2_reg_weight))
def __init__(self,
n_filters,
strides=1,
downsample=None,
regularization=0.01):
"""Initialize the BottleneckResidualUnit module.
Args:
n_filters: (int) the number of output filters.
strides: (int) the strides of the convolution.
downsample: a function to down-sample the feature maps.
regularization: L2 regularization factor for layer weights.
"""
super(BottleneckResidualUnit, self).__init__()
self.bn1 = BatchNormalization()
self.conv1 = Conv2D(n_filters,
1,
padding='same',
use_bias=False,
kernel_regularizer=regularizers.l2(regularization))
self.bn2 = BatchNormalization()
self.conv2 = Conv2D(n_filters,
3,
strides=strides,
padding='same',
use_bias=False,
kernel_regularizer=regularizers.l2(regularization))
self.bn3 = BatchNormalization()
self.conv3 = Conv2D(n_filters * self.expansion,
1,
padding='same',
use_bias=False,
kernel_regularizer=regularizers.l2(regularization))
self.leaky_relu = LeakyReLU()
self.downsample = downsample
def get_vggm_conv_block(x, conv_layers, use_bn, max_pool_size, name):
"""Conv block."""
l2_reg_weight = CONFIG.model.l2_reg_weight
for (channels, kernel_size) in conv_layers:
x = layers.Conv2D(channels,
kernel_size,
padding='same',
kernel_initializer='he_normal',
kernel_regularizer=regularizers.l2(l2_reg_weight),
bias_regularizer=regularizers.l2(l2_reg_weight))(x)
if use_bn:
x = layers.BatchNormalization()(x)
x = layers.Activation('relu')(x)
if max_pool_size > 0:
x = layers.MaxPooling2D(pool_size=max_pool_size,
strides=2,
padding='same',
name=name)(x)
else:
# Identity layer
x = layers.MaxPooling2D(pool_size=1,
strides=1,
padding='same',
name=name)(x)
return x
def get_conv_bn_layers(conv_params, use_bn, conv_dims=2):
"""Returns convolution and batch norm layers."""
if conv_dims == 1:
conv_layer = layers.Conv1D
elif conv_dims == 2:
conv_layer = layers.Conv2D
elif conv_dims == 3:
conv_layer = layers.Conv3D
else:
raise ValueError('Invalid number of conv_dims')
l2_reg_weight = CONFIG.MODEL.L2_REG_WEIGHT
conv_layers = []
bn_layers = []
activations = []
for channels, kernel_size, activate in conv_params:
if activate:
activation = tf.nn.relu
else:
activation = None
conv_layers.append(
conv_layer(
channels,
kernel_size,
padding='same',
kernel_regularizer=regularizers.l2(l2_reg_weight),
bias_regularizer=regularizers.l2(l2_reg_weight),
kernel_initializer='he_normal',
))
if use_bn:
bn_layers.append(layers.BatchNormalization())
activations.append(activation)
return conv_layers, bn_layers, activations
def rnn_model(params, training_dr_lstm=True, training_dr_ll=True):
"""RNN model for text."""
input_shape = (params['fix_len'])
seq_input = layers.Input(shape=input_shape)
# vocab+1 because of padding
seq_emb = layers.Embedding(params['vocab_size'] + 1,
params['emb_size'],
input_length=params['fix_len'])(seq_input)
lstm_out = layers.LSTM(params['hidden_lstm_size'],
dropout=params['dropout_rate_lstm'])(
seq_emb, training=training_dr_lstm)
out = layers.Dropout(rate=params['dropout_rate'],
seed=params['random_seed'])(lstm_out,
training=training_dr_ll)
if params['variational']:
# scale kl loss by number of training examples.
# larger training dataset depends less on prior
def scaled_kl_fn(p, q, _):
return tfp.distributions.kl_divergence(q, p) / params['n_train']
logits = tfpl.DenseReparameterization(
params['n_class_in'],
activation=None,
kernel_divergence_fn=scaled_kl_fn,
bias_posterior_fn=tfpl.util.default_mean_field_normal_fn(),
name='last_layer')(out)
else:
logits = layers.Dense(
params['n_class_in'],
activation=None,
kernel_regularizer=regularizers.l2(params['reg_weight']),
bias_regularizer=regularizers.l2(params['reg_weight']),
name='last_layer')(out)
probs = layers.Softmax(axis=1)(logits)
return models.Model(seq_input, probs, name='rnn')
def self_attn_block(inp, nc, squeeze_factor=8):
'''
Code borrows from https://github.com/taki0112/Self-Attention-GAN-Tensorflow
'''
assert nc // squeeze_factor > 0, f"Input channels must be >= {squeeze_factor}, recieved nc={nc}"
x = inp
shape_x = x.get_shape().as_list()
f = Conv2D(nc // squeeze_factor,
1,
kernel_regularizer=regularizers.l2(w_l2))(x)
g = Conv2D(nc // squeeze_factor,
1,
kernel_regularizer=regularizers.l2(w_l2))(x)
h = Conv2D(nc, 1, kernel_regularizer=regularizers.l2(w_l2))(x)
shape_f = f.get_shape().as_list()
shape_g = g.get_shape().as_list()
shape_h = h.get_shape().as_list()
flat_f = Reshape((-1, shape_f[-1]))(f)
flat_g = Reshape((-1, shape_g[-1]))(g)
flat_h = Reshape((-1, shape_h[-1]))(h)
s = Lambda(lambda x: K.batch_dot(x[0],
Permute((2, 1))(x[1])))([flat_g, flat_f])
beta = Softmax(axis=-1)(s)
o = Lambda(lambda x: K.batch_dot(x[0], x[1]))([beta, flat_h])
o = Reshape(shape_x[1:])(o)
o = Scale()(o)
out = add([o, inp])
return out
def SPADE_res_block(input_tensor,
cond_input_tensor,
f,
use_norm=True,
norm='none'):
"""
Semantic Image Synthesis with Spatially-Adaptive Normalization
Taesung Park, Ming-Yu Liu, Ting-Chun Wang, Jun-Yan Zhu
https://arxiv.org/abs/1903.07291
Note:
SPADE just works like a charm.
It speeds up training alot and is also a very promosing approach for solving profile face generation issue.
*(This implementation can be wrong since I haven't finished reading the paper.
The author hasn't release their code either (https://github.com/NVlabs/SPADE).)
"""
def SPADE(input_tensor, cond_input_tensor, f, use_norm=True, norm='none'):
x = input_tensor
x = normalization(x, norm, f) if use_norm else x
y = cond_input_tensor
y = Conv2D(128,
kernel_size=3,
kernel_regularizer=regularizers.l2(w_l2),
kernel_initializer=conv_init,
padding='same')(y)
y = Activation('relu')(y)
gamma = Conv2D(f,
kernel_size=3,
kernel_regularizer=regularizers.l2(w_l2),
kernel_initializer=conv_init,
padding='same')(y)
beta = Conv2D(f,
kernel_size=3,
kernel_regularizer=regularizers.l2(w_l2),
kernel_initializer=conv_init,
padding='same')(y)
x = add([x, multiply([x, gamma])])
x = add([x, beta])
return x
x = input_tensor
x = SPADE(x, cond_input_tensor, f, use_norm, norm)
x = Activation('relu')(x)
x = ReflectPadding2D(x)
x = Conv2D(f,
kernel_size=3,
kernel_regularizer=regularizers.l2(w_l2),
kernel_initializer=conv_init,
use_bias=not use_norm)(x)
x = SPADE(x, cond_input_tensor, f, use_norm, norm)
x = Activation('relu')(x)
x = ReflectPadding2D(x)
x = Conv2D(f,
kernel_size=3,
kernel_regularizer=regularizers.l2(w_l2),
kernel_initializer=conv_init)(x)
x = add([x, input_tensor])
x = Activation('relu')(x)
return x
def get_gru_layers(gru_params):
"""Returns GRU layers."""
l2_reg_weight = CONFIG.MODEL.L2_REG_WEIGHT
gru_layers = []
for units in gru_params:
gru_layers.append(
layers.CuDNNGRU(units=units,
kernel_regularizer=regularizers.l2(l2_reg_weight),
bias_regularizer=regularizers.l2(l2_reg_weight),
return_sequences=True))
return gru_layers
def residual_block(
inputs,
num_filters=16,
kernel_size=3,
strides=1,
activation="relu",
batch_normalization=True,
conv_first=True,
):
"""2D Convolution-Batch Normalization-Activation stack builder
# Arguments
inputs (tensor): input tensor from input image or previous layer
num_filters (int): Conv2D number of filters
kernel_size (int): Conv2D square kernel dimensions
strides (int): Conv2D square stride dimensions
activation (string): activation name
batch_normalization (bool): whether to include batch normalization
conv_first (bool): conv-bn-activation (True) or
bn-activation-conv (False)
# Returns
x (tensor): tensor as input to the next layer
"""
conv = Conv2D(
num_filters,
kernel_size=kernel_size,
strides=strides,
padding="same",
kernel_initializer="he_normal",
kernel_regularizer=l2(1e-4),
activation=None,
)
conv2 = Conv2D(
num_filters,
kernel_size=kernel_size,
strides=strides,
padding="same",
kernel_initializer="he_normal",
kernel_regularizer=l2(1e-4),
activation="linear",
)
x = conv(inputs)
x = BatchNormalization()(x)
x = Activation(activation)(x)
x = conv2(x)
x = add([inputs, x])
x = BatchNormalization()(x)
x = Activation(activation)(x)
return x
def get_fc_layers(fc_params):
"""Return fully connected layers."""
l2_reg_weight = CONFIG.MODEL.L2_REG_WEIGHT
fc_layers = []
for channels, activate in fc_params:
if activate:
activation = tf.nn.relu
else:
activation = None
fc_layers.append(
layers.Dense(channels, activation=activation,
kernel_regularizer=regularizers.l2(l2_reg_weight),
bias_regularizer=regularizers.l2(l2_reg_weight)))
return fc_layers
def normalization(inp, norm='none', group='16'):
x = inp
if norm == 'layernorm':
x = GroupNormalization(group=group)(x)
elif norm == 'batchnorm':
x = BatchNormalization()(x)
elif norm == 'groupnorm':
x = GroupNormalization(group=16)(x)
elif norm == 'instancenorm':
x = InstanceNormalization()(x)
elif norm == 'hybrid':
if group % 2 == 1:
raise ValueError(
f"Output channels must be an even number for hybrid norm, received {group}."
)
f = group
x0 = Lambda(lambda x: x[..., :f // 2])(x)
x1 = Lambda(lambda x: x[..., f // 2:])(x)
x0 = Conv2D(f // 2,
kernel_size=1,
kernel_regularizer=regularizers.l2(w_l2),
kernel_initializer=conv_init)(x0)
x1 = InstanceNormalization()(x1)
x = concatenate([x0, x1], axis=-1)
else:
x = x
return x
def bottleneck(filters, strides=(1, 1, 1), kernel_regularizer=l2(1e-4),
is_first_block_of_first_layer=False):
"""Basic 3 X 3 X 3 convolution blocks. Extended from raghakot's 2D impl."""
def f(input):
if is_first_block_of_first_layer:
# don't repeat bn->relu since we just did bn->relu->maxpool
conv_1_1 = Conv3D(filters=filters, kernel_size=(1, 1, 1),
strides=strides, padding="same",
kernel_initializer="he_normal",
kernel_regularizer=kernel_regularizer
)(input)
else:
conv_1_1 = _bn_relu_conv3d(filters=filters, kernel_size=(1, 1, 1),
strides=strides,
kernel_regularizer=kernel_regularizer
)(input)
conv_3_3 = _bn_relu_conv3d(filters=filters, kernel_size=(3, 3, 3),
kernel_regularizer=kernel_regularizer
)(conv_1_1)
residual = _bn_relu_conv3d(filters=filters * 4, kernel_size=(1, 1, 1),
kernel_regularizer=kernel_regularizer
)(conv_3_3)
return _shortcut3d(input, residual)
return f
def upscale_nn(input_tensor, f, use_norm=False, w_l2=w_l2, norm='none'):
x = input_tensor
x = UpSampling2D()(x)
x = ReflectPadding2D(x, 1)
x = Conv2D(f,
kernel_size=3,
kernel_regularizer=regularizers.l2(w_l2),
kernel_initializer=conv_init)(x)
x = normalization(x, norm, f) if use_norm else x
return x
def res_block(input_tensor, f, use_norm=False, w_l2=w_l2, norm='none'):
x = input_tensor
x = Conv2D(f,
kernel_size=3,
kernel_regularizer=regularizers.l2(w_l2),
kernel_initializer=conv_init,
use_bias=False,
padding="same")(x)
x = LeakyReLU(alpha=0.2)(x)
x = normalization(x, norm, f) if use_norm else x
x = Conv2D(f,
kernel_size=3,
kernel_regularizer=regularizers.l2(w_l2),
kernel_initializer=conv_init,
use_bias=False,
padding="same")(x)
x = add([x, input_tensor])
x = LeakyReLU(alpha=0.2)(x)
x = normalization(x, norm, f) if use_norm else x
return x
def upscale_ps(input_tensor, f, use_norm=False, w_l2=w_l2, norm='none'):
x = input_tensor
x = Conv2D(f * 4,
kernel_size=3,
kernel_regularizer=regularizers.l2(w_l2),
kernel_initializer=icnr_keras,
padding='same')(x)
x = LeakyReLU(0.2)(x)
x = normalization(x, norm, f) if use_norm else x
x = PixelShuffler()(x)
return x
def SPADE(input_tensor, cond_input_tensor, f, use_norm=True, norm='none'):
x = input_tensor
x = normalization(x, norm, f) if use_norm else x
y = cond_input_tensor
y = Conv2D(128,
kernel_size=3,
kernel_regularizer=regularizers.l2(w_l2),
kernel_initializer=conv_init,
padding='same')(y)
y = Activation('relu')(y)
gamma = Conv2D(f,
kernel_size=3,
kernel_regularizer=regularizers.l2(w_l2),
kernel_initializer=conv_init,
padding='same')(y)
beta = Conv2D(f,
kernel_size=3,
kernel_regularizer=regularizers.l2(w_l2),
kernel_initializer=conv_init,
padding='same')(y)
x = add([x, multiply([x, gamma])])
x = add([x, beta])
return x
def conv_block(input_tensor,
f,
use_norm=False,
strides=2,
w_l2=w_l2,
norm='none'):
x = input_tensor
x = Conv2D(f,
kernel_size=3,
strides=strides,
kernel_regularizer=regularizers.l2(w_l2),
kernel_initializer=conv_init,
use_bias=False,
padding="same")(x)
x = Activation("relu")(x)
x = normalization(x, norm, f) if use_norm else x
return x
def _conv_bn_relu3D(**conv_params):
filters = conv_params["filters"]
kernel_size = conv_params["kernel_size"]
strides = conv_params.setdefault("strides", (1, 1, 1))
kernel_initializer = conv_params.setdefault(
"kernel_initializer", "he_normal")
padding = conv_params.setdefault("padding", "same")
kernel_regularizer = conv_params.setdefault("kernel_regularizer",
l2(1e-4))
def f(input):
conv = Conv3D(filters=filters, kernel_size=kernel_size,
strides=strides, kernel_initializer=kernel_initializer,
padding=padding,
kernel_regularizer=kernel_regularizer)(input)
return _bn_relu(conv)
return f
def _bn_relu_conv3d(**conv_params):
"""Helper to build a BN -> relu -> conv3d block."""
filters = conv_params["filters"]
kernel_size = conv_params["kernel_size"]
strides = conv_params.setdefault("strides", (1, 1, 1))
kernel_initializer = conv_params.setdefault("kernel_initializer",
"he_normal")
padding = conv_params.setdefault("padding", "same")
kernel_regularizer = conv_params.setdefault("kernel_regularizer",
l2(1e-4))
def f(input):
activation = _bn_relu(input)
return Conv3D(filters=filters, kernel_size=kernel_size,
strides=strides, kernel_initializer=kernel_initializer,
padding=padding,
kernel_regularizer=kernel_regularizer)(activation)
return f
def build_graph(self):
self._construct_weights() #Construct p and q weights
saver, logits, KL = self.forward_pass()
log_softmax_var = tf.nn.log_softmax(logits) #Softmax the last layer
neg_ll = -tf.reduce_mean(
tf.reduce_sum(log_softmax_var * self.input_ph, axis=-1))
reg = l2(self.lam)
reg_var = reg(self.weights_q +
self.weights_p) # Apply regularization to weights
# tensorflow l2 regularization multiply 0.5 to the l2 norm multiply 2 so that it is back in the same scale
neg_ELBO = neg_ll + self.anneal_ph * KL + 2 * reg_var
train_op = tf.compat.v1.train.AdamOptimizer(self.lr).minimize(
neg_ELBO) #Train the model
return saver, logits, neg_ELBO, train_op
def _shortcut3d(input, residual):
"""3D shortcut to match input and residual and merges them with "sum"."""
stride_dim1 = ceil(input.get_shape().as_list()[DIM1_AXIS] \
/ residual.get_shape().as_list()[DIM1_AXIS])
stride_dim2 = ceil(input.get_shape().as_list()[DIM2_AXIS] \
/ residual.get_shape().as_list()[DIM2_AXIS])
stride_dim3 = ceil(input.get_shape().as_list()[DIM3_AXIS] \
/ residual.get_shape().as_list()[DIM3_AXIS])
equal_channels = residual.get_shape().as_list()[CHANNEL_AXIS] \
== input.get_shape().as_list()[CHANNEL_AXIS]
shortcut = input
if stride_dim1 > 1 or stride_dim2 > 1 or stride_dim3 > 1 \
or not equal_channels:
shortcut = Conv3D(
filters=residual.get_shape().as_list()[CHANNEL_AXIS],
kernel_size=(1, 1, 1),
strides=(stride_dim1, stride_dim2, stride_dim3),
kernel_initializer="he_normal", padding="valid",
kernel_regularizer=l2(1e-4)
)(input)
return add([shortcut, residual])
def __init__(self, regularization=0.01):
super(SiameseEncoder, self).__init__()
self.inplanes = 64
# Siamese branch.
self.siamese = Sequential([
Conv2D(64,
7,
strides=2,
padding='same',
use_bias=False,
kernel_regularizer=regularizers.l2(regularization)),
self._make_resblock(2,
128,
strides=2,
regularization=regularization),
self._make_resblock(2,
128,
strides=2,
regularization=regularization),
self._make_resblock(2,
256,
strides=2,
regularization=regularization),
])
# Merged main branch.
self.mainstream = Sequential([
self._make_resblock(2,
256,
strides=2,
regularization=regularization),
self._make_resblock(2,
256,
strides=2,
regularization=regularization),
])
self.bn = BatchNormalization()
self.leaky_relu = LeakyReLU()
def keras_build_fn(num_feature,
num_output,
is_sparse,
embedding_dim=-1,
num_hidden_layer=2,
hidden_layer_dim=512,
activation='elu',
learning_rate=1e-3,
dropout=0.5,
l1=0.0,
l2=0.0,
loss='categorical_crossentropy'):
"""Initializes and compiles a Keras DNN model using the Adam optimizer.
Args:
num_feature: number of features
num_output: number of outputs (targets, e.g., classes))
is_sparse: boolean whether input data is in sparse format
embedding_dim: int number of nodes in embedding layer; if value is <= 0 then
no embedding layer will be present in the model
num_hidden_layer: number of hidden layers
hidden_layer_dim: int number of nodes in the hidden layer(s)
activation: string
activation function for hidden layers; see https://keras.io/activations/
learning_rate: float learning rate for Adam
dropout: float proportion of nodes to dropout; values in [0, 1]
l1: float strength of L1 regularization on weights
l2: float strength of L2 regularization on weights
loss: string
loss function; see https://keras.io/losses/
Returns:
model: Keras.models.Model
compiled Keras model
"""
assert num_hidden_layer >= 1
inputs = Input(shape=(num_feature, ), sparse=is_sparse)
activation_func_args = ()
if activation.lower() == 'prelu':
activation_func = PReLU
elif activation.lower() == 'leakyrelu':
activation_func = LeakyReLU
elif activation.lower() == 'elu':
activation_func = ELU
elif activation.lower() == 'thresholdedrelu':
activation_func = ThresholdedReLU
else:
activation_func = Activation
activation_func_args = (activation)
if l1 > 0 and l2 > 0:
reg_init = lambda: regularizers.l1_l2(l1, l2)
elif l1 > 0:
reg_init = lambda: regularizers.l1(l1)
elif l2 > 0:
reg_init = lambda: regularizers.l2(l2)
else:
reg_init = lambda: None
if embedding_dim > 0:
# embedding layer
e = Dense(embedding_dim)(inputs)
x = Dense(hidden_layer_dim, kernel_regularizer=reg_init())(e)
x = activation_func(*activation_func_args)(x)
x = Dropout(dropout)(x)
else:
x = Dense(hidden_layer_dim, kernel_regularizer=reg_init())(inputs)
x = activation_func(*activation_func_args)(x)
x = Dropout(dropout)(x)
# add additional hidden layers
for _ in range(num_hidden_layer - 1):
x = Dense(hidden_layer_dim, kernel_regularizer=reg_init())(x)
x = activation_func(*activation_func_args)(x)
x = Dropout(dropout)(x)
x = Dense(num_output)(x)
preds = Activation('softmax')(x)
model = Model(inputs=inputs, outputs=preds)
model.compile(optimizer=Adam(lr=learning_rate), loss=loss)
return model
def identity_block_base(input_tensor,
kernel_size,
filters,
stage,
block,
num_updates,
dropout_rate=0.,
use_variational_layers=False):
"""The identity block is the block that has no conv layer at shortcut.
Arguments:
input_tensor: input tensor
kernel_size: default 3, the kernel size of
middle conv layer at main path
filters: list of integers, the filters of 3 conv layer at main path
stage: integer, current stage label, used for generating layer names
block: 'a','b'..., current block label, used for generating layer names
num_updates: integer, total steps in an epoch (for weighting the loss)
dropout_rate: float, always-on dropout rate.
use_variational_layers: boolean, if true train a variational model
Returns:
x: Output tensor for the block.
"""
filters1, filters2, filters3 = filters
divergence_fn = lambda q, p, ignore: (tfd.kl_divergence(q, p) / num_updates
)
if backend.image_data_format() == 'channels_last':
bn_axis = 3
else:
bn_axis = 1
conv_name_base = 'res' + str(stage) + block + '_branch'
bn_name_base = 'bn' + str(stage) + block + '_branch'
if not use_variational_layers:
first_conv_2d = layers.Conv2D(
filters1, (1, 1),
use_bias=False,
kernel_initializer='he_normal',
kernel_regularizer=regularizers.l2(L2_WEIGHT_DECAY),
name=conv_name_base + '2a')
if dropout_rate > 0.:
x = layers.Dropout(dropout_rate)(input_tensor, training=True)
x = first_conv_2d(x)
else:
x = first_conv_2d(input_tensor)
x = layers.BatchNormalization(axis=bn_axis,
momentum=BATCH_NORM_DECAY,
epsilon=BATCH_NORM_EPSILON,
name=bn_name_base + '2a')(x)
x = layers.Activation('relu')(x)
if dropout_rate > 0.:
x = layers.Dropout(dropout_rate)(x, training=True)
x = layers.Conv2D(filters2,
kernel_size,
use_bias=False,
padding='same',
kernel_initializer='he_normal',
kernel_regularizer=regularizers.l2(L2_WEIGHT_DECAY),
name=conv_name_base + '2b')(x)
x = layers.BatchNormalization(axis=bn_axis,
momentum=BATCH_NORM_DECAY,
epsilon=BATCH_NORM_EPSILON,
name=bn_name_base + '2b')(x)
x = layers.Activation('relu')(x)
if dropout_rate > 0.:
x = layers.Dropout(dropout_rate)(x, training=True)
x = layers.Conv2D(filters3, (1, 1),
use_bias=False,
kernel_initializer='he_normal',
kernel_regularizer=regularizers.l2(L2_WEIGHT_DECAY),
name=conv_name_base + '2c')(x)
x = layers.BatchNormalization(axis=bn_axis,
momentum=BATCH_NORM_DECAY,
epsilon=BATCH_NORM_EPSILON,
name=bn_name_base + '2c')(x)
else:
x = tfpl.Convolution2DFlipout(
filters1,
kernel_size=(1, 1),
padding='SAME',
name=conv_name_base + '2a',
kernel_divergence_fn=divergence_fn,
)(input_tensor)
x = layers.BatchNormalization(axis=bn_axis,
momentum=BATCH_NORM_DECAY,
epsilon=BATCH_NORM_EPSILON,
name=bn_name_base + '2a')(x)
x = layers.Activation('relu')(x)
x = tfpl.Convolution2DFlipout(
filters2,
kernel_size=kernel_size,
padding='SAME',
activation=None,
name=conv_name_base + '2b',
kernel_divergence_fn=divergence_fn,
)(x)
x = layers.BatchNormalization(axis=bn_axis,
momentum=BATCH_NORM_DECAY,
epsilon=BATCH_NORM_EPSILON,
name=bn_name_base + '2b')(x)
x = layers.Activation('relu')(x)
x = tfpl.Convolution2DFlipout(
filters3,
kernel_size=(1, 1),
padding='SAME',
activation=None,
name=conv_name_base + '2c',
kernel_divergence_fn=divergence_fn,
)(x)
x = layers.BatchNormalization(axis=bn_axis,
momentum=BATCH_NORM_DECAY,
epsilon=BATCH_NORM_EPSILON,
name=bn_name_base + '2c')(x)
x = layers.add([x, input_tensor])
x = layers.Activation('relu')(x)
return x
def ResNet50(method, num_classes, num_updates, dropout_rate):
"""Instantiates the ResNet50 architecture.
Args:
method: `str`, method for accounting for uncertainty. Must be one of
['vanilla', 'll_dropout', 'll_svi', 'dropout', 'svi', 'dropout_nofirst']
num_classes: `int` number of classes for image classification.
num_updates: integer, total steps in an epoch (for weighting the loss)
dropout_rate: Dropout rate for ll_dropout, dropout methods.
Returns:
A Keras model instance.
pylint: disable=invalid-name
"""
# Determine proper input shape
if backend.image_data_format() == 'channels_first':
input_shape = (3, 224, 224)
bn_axis = 1
else:
input_shape = (224, 224, 3)
bn_axis = 3
if (method in ['dropout', 'll_dropout', 'dropout_nofirst'
]) != (dropout_rate > 0.):
raise ValueError(
'Dropout rate should be nonzero iff a dropout method is used.'
'Method is {}, dropout is {}.'.format(method, dropout_rate))
use_variational_layers = method == 'svi'
hidden_layer_dropout = dropout_rate if method in [
'dropout', 'dropout_nofirst'
] else 0.
img_input = layers.Input(shape=input_shape)
x = layers.ZeroPadding2D(padding=(3, 3), name='conv1_pad')(img_input)
if (dropout_rate > 0.) and (method != 'dropout_nofirst'):
x = layers.Dropout(hidden_layer_dropout)(x, training=True)
x = layers.Conv2D(64, (7, 7),
use_bias=False,
strides=(2, 2),
padding='valid',
kernel_initializer='he_normal',
kernel_regularizer=regularizers.l2(L2_WEIGHT_DECAY),
name='conv1')(x)
x = layers.BatchNormalization(axis=bn_axis,
momentum=BATCH_NORM_DECAY,
epsilon=BATCH_NORM_EPSILON,
name='bn_conv1')(x)
x = layers.Activation('relu')(x)
x = layers.ZeroPadding2D(padding=(1, 1), name='pool1_pad')(x)
x = layers.MaxPooling2D((3, 3), strides=(2, 2))(x)
conv_block = functools.partial(
conv_block_base,
num_updates=num_updates,
dropout_rate=hidden_layer_dropout,
use_variational_layers=use_variational_layers)
identity_block = functools.partial(
identity_block_base,
num_updates=num_updates,
dropout_rate=hidden_layer_dropout,
use_variational_layers=use_variational_layers)
x = conv_block(x, 3, [64, 64, 256], stage=2, block='a', strides=(1, 1))
x = identity_block(x, 3, [64, 64, 256], stage=2, block='b')
x = identity_block(x, 3, [64, 64, 256], stage=2, block='c')
x = conv_block(x, 3, [128, 128, 512], stage=3, block='a')
x = identity_block(x, 3, [128, 128, 512], stage=3, block='b')
x = identity_block(x, 3, [128, 128, 512], stage=3, block='c')
x = identity_block(x, 3, [128, 128, 512], stage=3, block='d')
x = conv_block(x, 3, [256, 256, 1024], stage=4, block='a')
x = identity_block(x, 3, [256, 256, 1024], stage=4, block='b')
x = identity_block(x, 3, [256, 256, 1024], stage=4, block='c')
x = identity_block(x, 3, [256, 256, 1024], stage=4, block='d')
x = identity_block(x, 3, [256, 256, 1024], stage=4, block='e')
x = identity_block(x, 3, [256, 256, 1024], stage=4, block='f')
x = conv_block(x, 3, [512, 512, 2048], stage=5, block='a')
x = identity_block(x, 3, [512, 512, 2048], stage=5, block='b')
x = identity_block(x, 3, [512, 512, 2048], stage=5, block='c')
x = layers.GlobalAveragePooling2D(name='avg_pool')(x)
if dropout_rate > 0.:
x = layers.Dropout(dropout_rate)(x, training=True)
if method in ['ll_svi', 'svi']:
x = tfpl.dense_variational_v2.DenseVariational(
units=num_classes,
make_posterior_fn=posterior_mean_field,
make_prior_fn=functools.partial(prior_trainable,
num_updates=num_updates),
use_bias=True,
kl_weight=1. / num_updates,
kl_use_exact=True,
name='fc1000')(x)
else:
x = layers.Dense(num_classes,
kernel_regularizer=regularizers.l2(L2_WEIGHT_DECAY),
bias_regularizer=regularizers.l2(L2_WEIGHT_DECAY),
name='fc1000')(x)
# Create model.
return models.Model(img_input, x, name='resnet50')
def build(input_shape, num_outputs, block_fn, repetitions, reg_factor):
"""Instantiate a vanilla ResNet3D keras model.
# Arguments
input_shape: Tuple of input shape in the format
(conv_dim1, conv_dim2, conv_dim3, channels) if dim_ordering='tf'
(filter, conv_dim1, conv_dim2, conv_dim3) if dim_ordering='th'
num_outputs: The number of outputs at the final softmax layer
block_fn: Unit block to use {'basic_block', 'bottlenack_block'}
repetitions: Repetitions of unit blocks
# Returns
model: a 3D ResNet model that takes a 5D tensor (volumetric images
in batch) as input and returns a 1D vector (prediction) as output.
"""
_handle_data_format()
if len(input_shape) != 4:
raise ValueError("Input shape should be a tuple "
"(conv_dim1, conv_dim2, conv_dim3, channels) "
"for tensorflow as backend or "
"(channels, conv_dim1, conv_dim2, conv_dim3) "
"for theano as backend")
block_fn = _get_block(block_fn)
input = Input(shape=input_shape)
# first conv
conv1 = _conv_bn_relu3D(filters=64, kernel_size=(7, 7, 7),
strides=(2, 2, 2),
kernel_regularizer=l2(reg_factor)
)(input)
pool1 = MaxPooling3D(pool_size=(3, 3, 3), strides=(2, 2, 2),
padding="same")(conv1)
# repeat blocks
block = pool1
filters = 64
for i, r in enumerate(repetitions):
block = _residual_block3d(block_fn, filters=filters,
kernel_regularizer=l2(reg_factor),
repetitions=r, is_first_layer=(i == 0)
)(block)
filters *= 2
# last activation
block_output = _bn_relu(block)
# average poll and classification
pool2 = AveragePooling3D(pool_size=(block.get_shape().as_list()[DIM1_AXIS],
block.get_shape().as_list()[DIM2_AXIS],
block.get_shape().as_list()[DIM3_AXIS]),
strides=(1, 1, 1))(block_output)
flatten1 = Flatten()(pool2)
if num_outputs > 1:
dense = Dense(units=num_outputs,
kernel_initializer="he_normal",
activation="softmax",
kernel_regularizer=l2(reg_factor))(flatten1)
else:
dense = Dense(units=num_outputs,
kernel_initializer="he_normal",
activation="sigmoid",
kernel_regularizer=l2(reg_factor))(flatten1)
model = Model(inputs=input, outputs=dense)
return model
def stack_layers(inputs, net_layers, kernel_initializer='glorot_uniform'):
"""Builds the architecture of the network by applying each layer specified in net_layers to inputs.
Args:
inputs: a dict containing input_types and input_placeholders for each key
and value pair, respecively.
net_layers: a list of dicts containing all layers to be used in the
network, where each dict describes one such layer. each dict requires the
key 'type'. all other keys are dependent on the layer type.
kernel_initializer: initialization configuration passed to keras (see keras
initializers).
Returns:
outputs: a dict formatted in much the same way as inputs. it
contains input_types and output_tensors for each key and value pair,
respectively, where output_tensors are the outputs of the
input_placeholders in inputs after each layer in net_layers is applied.
"""
outputs = dict()
for key in inputs:
outputs[key] = inputs[key]
for layer in net_layers:
# check for l2_reg argument
l2_reg = layer.get('l2_reg')
if l2_reg:
l2_reg = l2(layer['l2_reg'])
# create the layer
if layer['type'] in [
'softplus', 'softsign', 'softmax', 'tanh', 'sigmoid', 'relu', 'selu'
]:
l = layers.Dense(
layer['size'],
activation=layer['type'],
kernel_initializer=kernel_initializer,
kernel_regularizer=l2_reg,
name=layer.get('name'))
elif layer['type'] == 'None':
l = layers.Dense(
layer['size'],
kernel_initializer=kernel_initializer,
kernel_regularizer=l2_reg,
name=layer.get('name'))
elif layer['type'] == 'Conv2D':
l = layers.Conv2D(
layer['channels'],
kernel_size=layer['kernel'],
activation='relu',
data_format='channels_last',
kernel_regularizer=l2_reg,
name=layer.get('name'))
elif layer['type'] == 'BatchNormalization':
l = layers.BatchNormalization(name=layer.get('name'))
elif layer['type'] == 'MaxPooling2D':
l = layers.MaxPooling2D(
pool_size=layer['pool_size'],
data_format='channels_first',
name=layer.get('name'))
elif layer['type'] == 'Dropout':
l = layers.Dropout(layer['rate'], name=layer.get('name'))
elif layer['type'] == 'Flatten':
l = layers.Flatten(name=layer.get('name'))
else:
raise ValueError("Invalid layer type '{}'".format(layer['type']))
# apply the layer to each input in inputs
for k in outputs:
outputs[k] = l(outputs[k])
return outputs
def stack_layers(inputs, layers, kernel_initializer='glorot_uniform'):
'''
Builds the architecture of the network by applying each layer specified in layers to inputs.
inputs: a dict containing input_types and input_placeholders for each key and value pair, respecively.
for spectralnet, this means the input_types 'Unlabeled' and 'Orthonorm'*
layers: a list of dicts containing all layers to be used in the network, where each dict describes
one such layer. each dict requires the key 'type'. all other keys are dependent on the layer
type
kernel_initializer: initialization configuration passed to keras (see keras initializers)
returns: outputs, a dict formatted in much the same way as inputs. it contains input_types and
output_tensors for each key and value pair, respectively, where output_tensors are
the outputs of the input_placeholders in inputs after each layer in layers is applied
* this is necessary since spectralnet takes multiple inputs and performs special computations on the
orthonorm layer
'''
outputs = dict()
for key in inputs:
outputs[key] = inputs[key]
for layer in layers:
# check for l2_reg argument
l2_reg = layer.get('l2_reg')
if l2_reg:
l2_reg = l2(layer['l2_reg'])
# create the layer
if layer['type'] == 'softplus_reg':
l = Dense(layer['size'],
activation='softplus',
kernel_initializer=kernel_initializer,
kernel_regularizer=l2(0.001),
name=layer.get('name'))
elif layer['type'] == 'softplus':
l = Dense(layer['size'],
activation='softplus',
kernel_initializer=kernel_initializer,
kernel_regularizer=l2_reg,
name=layer.get('name'))
elif layer['type'] == 'softmax':
l = Dense(layer['size'],
activation='softmax',
kernel_initializer=kernel_initializer,
kernel_regularizer=l2_reg,
name=layer.get('name'))
elif layer['type'] == 'tanh':
l = Dense(layer['size'],
activation='tanh',
kernel_initializer=kernel_initializer,
kernel_regularizer=l2_reg,
name=layer.get('name'))
elif layer['type'] == 'relu':
l = Dense(layer['size'],
activation='relu',
kernel_initializer=kernel_initializer,
kernel_regularizer=l2_reg,
name=layer.get('name'))
elif layer['type'] == 'selu':
l = Dense(layer['size'],
activation='selu',
kernel_initializer=kernel_initializer,
kernel_regularizer=l2_reg,
name=layer.get('name'))
elif layer['type'] == 'Conv2D':
l = Conv2D(layer['channels'],
kernel_size=layer['kernel'],
activation='relu',
data_format='channels_last',
kernel_regularizer=l2_reg,
name=layer.get('name'))
elif layer['type'] == 'BatchNormalization':
l = BatchNormalization(name=layer.get('name'))
elif layer['type'] == 'MaxPooling2D':
l = MaxPooling2D(pool_size=layer['pool_size'],
data_format='channels_first',
name=layer.get('name'))
elif layer['type'] == 'Dropout':
l = Dropout(layer['rate'], name=layer.get('name'))
elif layer['type'] == 'Flatten':
l = Flatten(name=layer.get('name'))
elif layer['type'] == 'Orthonorm':
l = Orthonorm(outputs['Orthonorm'], name=layer.get('name'))
else:
raise ValueError("Invalid layer type '{}'".format(layer['type']))
# apply the layer to each input in inputs
for k in outputs:
outputs[k] = l(outputs[k])
return outputs
def dual_attn_block(inp, nc, squeeze_factor=8):
'''
https://github.com/junfu1115/DANet
'''
assert nc // squeeze_factor > 0, f"Input channels must be >= {squeeze_factor}, recieved nc={nc}"
x = inp
shape_x = x.get_shape().as_list()
# position attention module
x_pam = Conv2D(nc,
kernel_size=3,
kernel_regularizer=regularizers.l2(w_l2),
kernel_initializer=conv_init,
use_bias=False,
padding="same")(x)
x_pam = Activation("relu")(x_pam)
x_pam = normalization(x_pam, norm, nc)
f_pam = Conv2D(nc // squeeze_factor,
1,
kernel_regularizer=regularizers.l2(w_l2))(x_pam)
g_pam = Conv2D(nc // squeeze_factor,
1,
kernel_regularizer=regularizers.l2(w_l2))(x_pam)
h_pam = Conv2D(nc, 1, kernel_regularizer=regularizers.l2(w_l2))(x_pam)
shape_f_pam = f_pam.get_shape().as_list()
shape_g_pam = g_pam.get_shape().as_list()
shape_h_pam = h_pam.get_shape().as_list()
flat_f_pam = Reshape((-1, shape_f_pam[-1]))(f_pam)
flat_g_pam = Reshape((-1, shape_g_pam[-1]))(g_pam)
flat_h_pam = Reshape((-1, shape_h_pam[-1]))(h_pam)
s_pam = Lambda(lambda x: K.batch_dot(x[0],
Permute((2, 1))(x[1])))(
[flat_g_pam, flat_f_pam])
beta_pam = Softmax(axis=-1)(s_pam)
o_pam = Lambda(lambda x: K.batch_dot(x[0], x[1]))([beta_pam, flat_h_pam])
o_pam = Reshape(shape_x[1:])(o_pam)
o_pam = Scale()(o_pam)
out_pam = add([o_pam, x_pam])
out_pam = Conv2D(nc,
kernel_size=3,
kernel_regularizer=regularizers.l2(w_l2),
kernel_initializer=conv_init,
use_bias=False,
padding="same")(out_pam)
out_pam = Activation("relu")(out_pam)
out_pam = normalization(out_pam, norm, nc)
# channel attention module
x_chn = Conv2D(nc,
kernel_size=3,
kernel_regularizer=regularizers.l2(w_l2),
kernel_initializer=conv_init,
use_bias=False,
padding="same")(x)
x_chn = Activation("relu")(x_chn)
x_chn = normalization(x_chn, norm, nc)
shape_x_chn = x_chn.get_shape().as_list()
flat_f_chn = Reshape((-1, shape_x_chn[-1]))(x_chn)
flat_g_chn = Reshape((-1, shape_x_chn[-1]))(x_chn)
flat_h_chn = Reshape((-1, shape_x_chn[-1]))(x_chn)
s_chn = Lambda(lambda x: K.batch_dot(Permute((2, 1))(x[0]), x[1]))(
[flat_g_chn, flat_f_chn])
s_new_chn = Lambda(lambda x: K.repeat_elements(K.max(x, -1, keepdims=True),
nc, -1))(s_chn)
s_new_chn = Lambda(lambda x: x[0] - x[1])([s_new_chn, s_chn])
beta_chn = Softmax(axis=-1)(s_new_chn)
o_chn = Lambda(lambda x: K.batch_dot(x[0],
Permute((2, 1))(x[1])))(
[flat_h_chn, beta_chn])
o_chn = Reshape(shape_x[1:])(o_chn)
o_chn = Scale()(o_chn)
out_chn = add([o_chn, x_chn])
out_chn = Conv2D(nc,
kernel_size=3,
kernel_regularizer=regularizers.l2(w_l2),
kernel_initializer=conv_init,
use_bias=False,
padding="same")(out_chn)
out_chn = Activation("relu")(out_chn)
out_chn = normalization(out_chn, norm, nc)
out = add([out_pam, out_chn])
return out
def __init__(self, n_out, regularization=0.01):
"""Initialize the DirectionNet.
Args:
n_out: (int) the number of output distributions.
regularization: L2 regularization factor for layer weights.
"""
super(DirectionNet, self).__init__()
self.encoder = SiameseEncoder()
self.inplanes = self.encoder.inplanes
self.decoder_block1 = Sequential([
Conv2D(256,
3,
use_bias=False,
kernel_regularizer=regularizers.l2(regularization)),
self._make_resblock(2, 128, regularization=regularization),
BatchNormalization(),
LeakyReLU()
])
self.decoder_block2 = Sequential([
Conv2D(128,
3,
use_bias=False,
kernel_regularizer=regularizers.l2(regularization)),
self._make_resblock(2, 64, regularization=regularization),
BatchNormalization(),
LeakyReLU()
])
self.decoder_block3 = Sequential([
Conv2D(64,
3,
use_bias=False,
kernel_regularizer=regularizers.l2(regularization)),
self._make_resblock(2, 32, regularization=regularization),
BatchNormalization(),
LeakyReLU()
])
self.decoder_block4 = Sequential([
Conv2D(32,
3,
use_bias=False,
kernel_regularizer=regularizers.l2(regularization)),
self._make_resblock(2, 16, regularization=regularization),
BatchNormalization(),
LeakyReLU()
])
self.decoder_block5 = Sequential([
Conv2D(16,
3,
use_bias=False,
kernel_regularizer=regularizers.l2(regularization)),
self._make_resblock(2, 8, regularization=regularization),
BatchNormalization(),
LeakyReLU()
])
self.decoder_block6 = Sequential([
Conv2D(8,
3,
use_bias=False,
kernel_regularizer=regularizers.l2(regularization)),
self._make_resblock(2, 4, regularization=regularization),
BatchNormalization(),
LeakyReLU()
])
self.down_channel = Conv2D(
n_out, 1, kernel_regularizer=regularizers.l2(regularization))