def resnet20(
inputs: tf_compat.Tensor,
training: Union[bool, tf_compat.Tensor] = True,
num_classes: int = 10,
class_type: str = "single",
kernel_initializer=tf_compat.glorot_uniform_initializer(),
bias_initializer=tf_compat.zeros_initializer(),
beta_initializer=tf_compat.zeros_initializer(),
gamma_initializer=tf_compat.ones_initializer(),
) -> tf_compat.Tensor:
with tf_compat.variable_scope("resnet20", reuse=tf_compat.AUTO_REUSE):
sec_settings = [
ResNetSection(num_blocks=2, out_channels=16, downsample=False),
ResNetSection(num_blocks=2, out_channels=32, downsample=True),
ResNetSection(num_blocks=2, out_channels=64, downsample=True),
]
net = resnet_const(
inputs,
training,
sec_settings,
num_classes,
class_type=class_type,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
beta_initializer=beta_initializer,
gamma_initializer=gamma_initializer,
simplified_arch=True,
)
return net
def _basic_block(
name: str,
x_tens: tf_compat.Tensor,
training: Union[bool, tf_compat.Tensor],
out_channels: int,
stride: int,
kernel_initializer,
bias_initializer,
beta_initializer,
gamma_initializer,
) -> tf_compat.Tensor:
with tf_compat.variable_scope(name, reuse=tf_compat.AUTO_REUSE):
out = conv2d_block(
"conv_bn_0",
x_tens,
training,
out_channels,
kernel_size=3,
stride=stride,
padding=1,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
beta_initializer=beta_initializer,
gamma_initializer=gamma_initializer,
)
out = conv2d_block(
"conv_bn_1",
out,
training,
out_channels,
kernel_size=3,
padding=1,
act=None,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
beta_initializer=beta_initializer,
gamma_initializer=gamma_initializer,
)
if stride > 1 or int(x_tens.shape[3]) != out_channels:
out = tf_compat.add(
out,
_identity_modifier(
x_tens,
training,
out_channels,
stride,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
beta_initializer=beta_initializer,
gamma_initializer=gamma_initializer,
),
)
else:
out = tf_compat.add(out, x_tens)
out = activation(out, act="relu", name="act_out")
return out
def create(
self,
name: str,
x_tens: tf_compat.Tensor,
training: Union[bool, tf_compat.Tensor],
kernel_initializer,
bias_initializer,
beta_initializer,
gamma_initializer,
) -> tf_compat.Tensor:
"""
Create the section in the current graph and scope
:param name: the name for the scope to create the section under
:param x_tens: The input tensor to the ResNet architecture
:param training: bool or Tensor to specify if the model should be run
in training or inference mode
:param kernel_initializer: Initializer to use for the conv and
fully connected kernels
:param bias_initializer: Initializer to use for the bias in the fully connected
:param beta_initializer: Initializer to use for the batch norm beta variables
:param gamma_initializer: Initializer to use for the batch norm gama variables
:return: the output tensor from the section
"""
out = x_tens
with tf_compat.variable_scope(name, reuse=tf_compat.AUTO_REUSE):
stride = 2 if self.downsample else 1
for block in range(self.num_blocks):
if self.proj_channels > 0:
out = _bottleneck_block(
name="block_{}".format(block),
x_tens=out,
training=training,
out_channels=self.out_channels,
proj_channels=self.proj_channels,
stride=stride,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
beta_initializer=beta_initializer,
gamma_initializer=gamma_initializer,
)
else:
out = _basic_block(
name="block_{}".format(block),
x_tens=out,
training=training,
out_channels=self.out_channels,
stride=stride,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
beta_initializer=beta_initializer,
gamma_initializer=gamma_initializer,
)
stride = 1
return out
def pool2d(
name: str,
x_tens: tf_compat.Tensor,
type_: str,
pool_size: Union[int, Tuple[int, int]],
strides: Union[int, Tuple[int, int]] = 1,
padding: Union[str, int, Tuple[int, ...]] = "same",
data_format: str = "channels_last",
):
"""
Create a pool op with the given name in the current graph and scope.
Supported are [max, avg, global_avg]
:param name: the name to given to the pooling op in the graph
:param x_tens: the input tensor to apply pooling to
:param type_: the type of pooling to apply, one of [max, avg, global_avg]
:param pool_size: the size of the pooling window to apply,
if global_avg then is the desired output size
:param strides: the stride to apply for the pooling op,
if global_avg then is unused
:param padding: any padding to apply to the tensor before pooling;
if string then uses tensorflows built in padding, else uses symmetric_pad2d
:param data_format: either channels_last or channels_first
:return: the tensor after pooling
"""
with tf_compat.variable_scope(name, reuse=tf_compat.AUTO_REUSE):
out = symmetric_pad2d(x_tens, padding, data_format)
if type_ == "max":
return tf_compat.layers.max_pooling2d(
out,
pool_size,
strides,
padding if isinstance(padding, str) else "valid",
data_format,
)
elif type_ == "avg":
return tf_compat.layers.average_pooling2d(
out,
pool_size,
strides,
padding if isinstance(padding, str) else "valid",
data_format,
)
elif type_ == "global_avg":
if pool_size != 1 and pool_size != (1, 1):
raise ValueError(
"only output pool_size of 1 is supported for global average pooling"
)
return tf_compat.reduce_mean(
out,
[1, 2] if data_format == "channels_last" else [2, 3],
keepdims=True,
)
else:
raise ValueError("unrecognized type_ given of {}".format(type_))
def _inverted_bottleneck_block(
name: str,
x_tens: tf_compat.Tensor,
training: Union[bool, tf_compat.Tensor],
out_channels: int,
exp_channels: int,
stride: int,
kernel_initializer,
bias_initializer,
beta_initializer,
gamma_initializer,
) -> tf_compat.Tensor:
with tf_compat.variable_scope(name, reuse=tf_compat.AUTO_REUSE):
out = conv2d_block(
"expand",
x_tens,
training,
exp_channels,
kernel_size=1,
act="relu6",
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
beta_initializer=beta_initializer,
gamma_initializer=gamma_initializer,
)
out = depthwise_conv2d_block(
"spatial",
out,
training,
exp_channels,
kernel_size=3,
stride=stride,
act="relu6",
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
beta_initializer=beta_initializer,
gamma_initializer=gamma_initializer,
)
out = conv2d_block(
"compress",
out,
training,
out_channels,
kernel_size=1,
act=None,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
beta_initializer=beta_initializer,
gamma_initializer=gamma_initializer,
)
if stride == 1 and int(x_tens.shape[3]) == out_channels:
out = tf_compat.add(out, x_tens)
return out
def mobilenet_const(
x_tens: tf_compat.Tensor,
training: Union[bool, tf_compat.Tensor],
sec_settings: List[MobileNetSection],
num_classes: int,
class_type: str,
kernel_initializer,
bias_initializer,
beta_initializer,
gamma_initializer,
) -> tf_compat.Tensor:
"""
Graph constructor for MobileNet implementation.
:param x_tens: The input tensor to the MobileNet architecture
:param training: bool or Tensor to specify if the model should be run
in training or inference mode
:param sec_settings: The settings for each section in the MobileNet modoel
:param num_classes: The number of classes to classify
:param class_type: One of [single, multi, None] to support multi class training.
Default single. If None, then will not add the fully connected at the end.
:param kernel_initializer: Initializer to use for the conv and
fully connected kernels
:param bias_initializer: Initializer to use for the bias in the fully connected
:param beta_initializer: Initializer to use for the batch norm beta variables
:param gamma_initializer: Initializer to use for the batch norm gama variables
:return: the output tensor from the created graph
"""
with tf_compat.variable_scope(BASE_NAME_SCOPE, reuse=tf_compat.AUTO_REUSE):
out = _input(x_tens, training, kernel_initializer, bias_initializer,
gamma_initializer)
for sec_index, section in enumerate(sec_settings):
out = section.create(
name="section_{}".format(sec_index),
x_tens=out,
training=training,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
beta_initializer=beta_initializer,
gamma_initializer=gamma_initializer,
)
logits = _classifier(
out,
training,
num_classes,
class_type,
kernel_initializer,
bias_initializer,
beta_initializer,
gamma_initializer,
)
return logits
def create_op_pruning_no_update(
op: tf_compat.Operation,
op_input: tf_compat.Tensor,
ks_group: str,
leave_enabled: bool = True,
is_after_end_step: tf_compat.Tensor = None,
) -> PruningOpVars:
"""
Creates the necessary variables and operators to gradually
apply sparsity to an operators variable without returning a
PruningOpVars.update value.
:param op: the operation to prune to the given sparsity
:param op_input: the parameter within the op to create a mask for
:param ks_group: the group identifier the scope should be created under
mask_creator
:param leave_enabled: True to continue masking the weights after end_epoch,
False to stop masking
:param is_after_end_step: only should be provided if leave_enabled is False;
tensor that is true if the current global step is after end_epoch
:return: a named tuple containing the assignment op, mask variable,
threshold tensor, and masked tensor
"""
if tf_contrib_err:
raise tf_contrib_err
op_sgv = graph_editor.sgv(op)
# create the necessary variables first
with tf_compat.variable_scope(PruningScope.model(op, ks_group),
reuse=tf_compat.AUTO_REUSE):
mask = tf_compat.get_variable(
PruningScope.VAR_MASK,
op_input.get_shape(),
initializer=tf_compat.ones_initializer(),
trainable=False,
dtype=op_input.dtype,
)
tf_compat.add_to_collection(
PruningScope.collection_name(ks_group, PruningScope.VAR_MASK), mask)
# create the masked operation and assign as the new input to the op
with tf_compat.name_scope(
PruningScope.model(op, ks_group, trailing_slash=True)):
masked = tf_compat.multiply(mask, op_input, PruningScope.OP_MASKED_VAR)
op_inp_tens = (masked if leave_enabled else tf_compat.cond(
is_after_end_step, lambda: op_input, lambda: masked))
op_swapped_inputs = [
inp if inp != op_input else op_inp_tens for inp in op_sgv.inputs
]
graph_editor.swap_inputs(op, op_swapped_inputs)
tf_compat.add_to_collection(
PruningScope.collection_name(ks_group, PruningScope.OP_MASKED_VAR),
masked)
return PruningOpVars(op, op_input, None, mask, masked)
def create(
self,
name: str,
x_tens: tf_compat.Tensor,
training: Union[bool, tf_compat.Tensor],
kernel_initializer,
bias_initializer,
beta_initializer,
gamma_initializer,
) -> tf_compat.Tensor:
"""
Create the section in the current graph and scope
:param name: the name for the scope to create the section under
:param x_tens: The input tensor to the MobileNet architecture
:param training: bool or Tensor to specify if the model should be run
in training or inference mode
:param kernel_initializer: Initializer to use for the conv and
fully connected kernels
:param bias_initializer: Initializer to use for the bias in the fully connected
:param beta_initializer: Initializer to use for the batch norm beta variables
:param gamma_initializer: Initializer to use for the batch norm gama variables
:return: the output tensor from the section
"""
out = x_tens
with tf_compat.variable_scope(name, reuse=tf_compat.AUTO_REUSE):
for block in range(self.num_blocks):
out = conv2d_block(
name="block_{}".format(block),
x_tens=out,
training=training,
channels=self.out_channels,
kernel_size=3,
include_bn=self.use_batchnorm,
include_bias=True,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
beta_initializer=beta_initializer,
gamma_initializer=gamma_initializer,
)
out = pool2d(
name="pool",
x_tens=out,
type_="max",
pool_size=2,
strides=2,
padding="valid",
)
return out
def conv2d(
name: str,
x_tens: tf_compat.Tensor,
in_chan: int,
out_chan: int,
kernel: int,
stride: int,
padding: str,
act: Union[None, str] = None,
):
"""
Create a convolutional layer with the proper ops and variables.
:param name: the name scope to create the layer under
:param x_tens: the tensor to apply the layer to
:param in_chan: the number of input channels
:param out_chan: the number of output channels
:param kernel: the kernel size to create a convolution for
:param stride: the stride to apply to the convolution
:param padding: the padding to apply to the convolution
:param act: an activation type to add into the layer, supported:
[None, relu, sigmoid, softmax]
:return: the created layer
"""
with tf_compat.variable_scope(name, reuse=tf_compat.AUTO_REUSE):
weight = tf_compat.get_variable(
"weight",
shape=[kernel, kernel, in_chan, out_chan],
initializer=tf_compat.glorot_normal_initializer(),
dtype=tf_compat.float32,
)
bias = tf_compat.get_variable(
"bias",
shape=[out_chan],
initializer=tf_compat.zeros_initializer(),
dtype=tf_compat.float32,
)
x_tens = tf_compat.nn.conv2d(x_tens,
weight,
strides=[1, stride, stride, 1],
padding=padding,
name="conv")
x_tens = tf_compat.nn.bias_add(x_tens, bias, name="bias_add")
x_tens = activation(x_tens, act)
return x_tens
def _classifier(
x_tens: tf_compat.Tensor,
training: Union[bool, tf_compat.Tensor],
num_classes: int,
class_type: str,
kernel_initializer,
bias_initializer,
beta_initializer,
gamma_initializer,
) -> tf_compat.Tensor:
with tf_compat.variable_scope("classifier", reuse=tf_compat.AUTO_REUSE):
logits = pool2d(name="avgpool", x_tens=x_tens, type_="global_avg", pool_size=1)
if num_classes:
logits = tf_compat.layers.dropout(
logits, 0.2, training=training, name="dropout"
)
logits = tf_compat.reshape(logits, [-1, int(logits.shape[3])])
if class_type:
if class_type == "single":
act = "softmax"
elif class_type == "multi":
act = "sigmoid"
else:
raise ValueError(
"unknown class_type given of {}".format(class_type)
)
else:
act = None
logits = dense_block(
"dense",
logits,
training,
num_classes,
include_bn=False,
act=act,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
beta_initializer=beta_initializer,
gamma_initializer=gamma_initializer,
)
return logits
def fc(
name: str,
x_tens: tf_compat.Tensor,
in_chan: int,
out_chan: int,
act: Union[None, str] = None,
):
"""
Create a fully connected layer with the proper ops and variables.
:param name: the name scope to create the layer under
:param x_tens: the tensor to apply the layer to
:param in_chan: the number of input channels
:param out_chan: the number of output channels
:param act: an activation type to add into the layer, supported:
[None, relu, sigmoid, softmax]
:return: the created layer
"""
with tf_compat.variable_scope(name, reuse=tf_compat.AUTO_REUSE):
weight = tf_compat.get_variable(
"weight",
shape=[in_chan, out_chan],
initializer=tf_compat.glorot_normal_initializer(),
dtype=tf_compat.float32,
)
bias = tf_compat.get_variable(
"bias",
shape=[out_chan],
initializer=tf_compat.zeros_initializer(),
dtype=tf_compat.float32,
)
x_tens = tf_compat.matmul(x_tens, weight, name="matmul")
x_tens = tf_compat.nn.bias_add(x_tens, bias, name="bias_add")
x_tens = activation(x_tens, act)
return x_tens
def _dw_sep_block(
name: str,
x_tens: tf_compat.Tensor,
training: Union[bool, tf_compat.Tensor],
out_channels: int,
stride: int,
kernel_initializer,
bias_initializer,
beta_initializer,
gamma_initializer,
) -> tf_compat.Tensor:
with tf_compat.variable_scope(name, reuse=tf_compat.AUTO_REUSE):
out = depthwise_conv2d_block(
"depth",
x_tens,
training,
int(x_tens.shape[3]),
kernel_size=3,
padding=1,
stride=stride,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
gamma_initializer=gamma_initializer,
)
out = conv2d_block(
"point",
out,
training,
out_channels,
kernel_size=1,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
beta_initializer=beta_initializer,
gamma_initializer=gamma_initializer,
)
return out
def mnist_net(inputs: tf_compat.Tensor,
num_classes: int = 10,
act: str = None) -> tf_compat.Tensor:
"""
A simple convolutional model created for the MNIST dataset
:param inputs: the inputs tensor to create the network for
:param num_classes: the number of classes to create the final layer for
:param act: the final activation to use in the model,
supported: [None, relu, sigmoid, softmax]
:return: the logits output from the created network
"""
if act not in [None, "sigmoid", "softmax"]:
raise ValueError("unsupported value for act given of {}".format(act))
with tf_compat.variable_scope(BASE_NAME_SCOPE, reuse=tf_compat.AUTO_REUSE):
with tf_compat.variable_scope("blocks", reuse=tf_compat.AUTO_REUSE):
x_tens = conv2d(
name="conv0",
x_tens=inputs,
in_chan=1,
out_chan=16,
kernel=5,
stride=1,
padding="SAME",
act="relu",
)
x_tens = conv2d(
name="conv1",
x_tens=x_tens,
in_chan=16,
out_chan=32,
kernel=5,
stride=2,
padding="SAME",
act="relu",
)
x_tens = conv2d(
name="conv2",
x_tens=x_tens,
in_chan=32,
out_chan=64,
kernel=5,
stride=1,
padding="SAME",
act="relu",
)
x_tens = conv2d(
name="conv3",
x_tens=x_tens,
in_chan=64,
out_chan=128,
kernel=5,
stride=2,
padding="SAME",
act="relu",
)
with tf_compat.variable_scope("classifier"):
x_tens = tf_compat.reduce_mean(x_tens, axis=[1, 2])
x_tens = tf_compat.reshape(x_tens, [-1, 128])
x_tens = fc(name="fc",
x_tens=x_tens,
in_chan=128,
out_chan=num_classes)
with tf_compat.variable_scope("logits"):
logits = activation(x_tens, act)
return logits
def create(
self,
name: str,
x_tens: tf_compat.Tensor,
training: Union[bool, tf_compat.Tensor],
kernel_initializer,
bias_initializer,
beta_initializer,
gamma_initializer,
) -> tf_compat.Tensor:
"""
Create the section in the current graph and scope
:param name: the name for the scope to create the section under
:param x_tens: The input tensor to the MobileNet architecture
:param training: bool or Tensor to specify if the model should be run
in training or inference mode
:param kernel_initializer: Initializer to use for the conv and
fully connected kernels
:param bias_initializer: Initializer to use for the bias in the fully connected
:param beta_initializer: Initializer to use for the batch norm beta variables
:param gamma_initializer: Initializer to use for the batch norm gama variables
:return: the output tensor from the section
"""
out = x_tens
with tf_compat.variable_scope(name, reuse=tf_compat.AUTO_REUSE):
stride = 2 if self.downsample else 1
exp_channels = (
self.exp_channels
if self.exp_channels is not None
else _make_divisible(int(out.shape[3]) * self.exp_ratio, 8)
)
for block in range(self.num_blocks):
if self.init_section and block == 0:
out = _input_inverted_bottleneck_block(
name="block_{}".format(block),
x_tens=out,
training=training,
out_channels=self.out_channels,
exp_channels=exp_channels,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
beta_initializer=beta_initializer,
gamma_initializer=gamma_initializer,
)
else:
out = _inverted_bottleneck_block(
name="block_{}".format(block),
x_tens=out,
training=training,
out_channels=self.out_channels,
exp_channels=exp_channels,
stride=stride,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
beta_initializer=beta_initializer,
gamma_initializer=gamma_initializer,
)
stride = 1
exp_channels = (
self.exp_channels
if self.exp_channels is not None
else _make_divisible(self.out_channels * self.exp_ratio, 8)
)
return out
def depthwise_conv2d_block(
name: str,
x_tens: tf_compat.Tensor,
training: Union[bool, tf_compat.Tensor],
channels: int,
kernel_size: int,
padding: Union[str, int, Tuple[int, ...]] = "same",
stride: int = 1,
data_format: str = "channels_last",
include_bn: bool = True,
include_bias: bool = None,
act: Union[None, str] = "relu",
kernel_initializer=tf_compat.glorot_uniform_initializer(),
bias_initializer=tf_compat.zeros_initializer(),
beta_initializer=tf_compat.zeros_initializer(),
gamma_initializer=tf_compat.ones_initializer(),
):
"""
Create a depthwise convolution op and supporting ops (batch norm, activation, etc)
in the current graph and scope.
:param name: The name to group all ops under in the graph
:param x_tens: The input tensor to apply a convolution and supporting ops to
:param training: A bool or tensor to indicate if the net is being run
in training mode or not. Used for batch norm
:param channels: The number of output channels from the conv op
:param kernel_size: The size of the kernel to use for the conv op
:param padding: Any padding to apply to the tensor before the convolution;
if string then uses tensorflows built in padding, else uses symmetric_pad2d
:param stride: The stride to apply for the convolution
:param data_format: Either channels_last or channels_first
:param include_bn: True to include a batch norm operation after the conv,
False otherwise
:param include_bias: If left unset, will add a bias if not include_bn.
Otherwise can be set to True to include a bias after the convolution,
False otherwise.
:param act: The activation to apply after the conv op and batch norm (if included).
Default is "relu", set to None for no activation.
:param kernel_initializer: The initializer to use for the convolution kernels
:param bias_initializer: The initializer to use for the bias variable,
if a bias is included
:param beta_initializer: The initializer to use for the beta variable,
if batch norm is included
:param gamma_initializer: The initializer to use for the gamma variable,
if gamma is included
:return: the tensor after all ops have been applied
"""
if include_bias is None:
include_bias = not include_bn
channel_axis = 3 if data_format == "channels_last" else 1
stride = ([1, stride, stride, 1]
if data_format == "channels_last" else [1, stride, stride, 1])
kernel_shape = (kernel_size, kernel_size, int(x_tens.shape[channel_axis]),
1)
with tf_compat.variable_scope(name, reuse=tf_compat.AUTO_REUSE):
with tf_compat.variable_scope("conv"):
kernel = tf_compat.get_variable(
"kernel",
shape=kernel_shape,
initializer=kernel_initializer,
trainable=True,
)
bias = (tf_compat.get_variable(
"bias",
shape=(channels, ),
initializer=bias_initializer,
trainable=True,
) if include_bias else None)
out = symmetric_pad2d(x_tens, padding, data_format)
out = tf_compat.nn.depthwise_conv2d(
out,
kernel,
stride,
padding=padding.upper()
if isinstance(padding, str) else "VALID",
data_format="NHWC"
if data_format == "channels_last" else "NCHW",
)
if bias is not None:
out = tf_compat.nn.bias_add(out, bias, data_format)
if include_bn:
out = tf_compat.layers.batch_normalization(
out,
axis=3 if data_format == "channels_last" else 1,
momentum=BN_MOMENTUM,
epsilon=BN_EPSILON,
beta_initializer=beta_initializer,
gamma_initializer=gamma_initializer,
training=training,
name="bn",
)
out = activation(out, act)
return out
def _classifier(
x_tens: tf_compat.Tensor,
training: Union[bool, tf_compat.Tensor],
num_classes: int,
class_type: str,
kernel_initializer,
bias_initializer,
beta_initializer,
gamma_initializer,
) -> tf_compat.Tensor:
with tf_compat.variable_scope("classifier", reuse=tf_compat.AUTO_REUSE):
if num_classes:
if class_type:
if class_type == "single":
final_act = "softmax"
elif class_type == "multi":
final_act = "sigmoid"
else:
raise ValueError(
"unknown class_type given of {}".format(class_type))
else:
final_act = None
out = tf_compat.transpose(x_tens, [0, 3, 1, 2])
out = tf_compat.reshape(out, [-1, 7 * 7 * 512])
out = dense_block(
"mlp_0",
out,
training,
channels=4096,
include_bn=False,
include_bias=True,
dropout_rate=0.5,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
beta_initializer=beta_initializer,
gamma_initializer=gamma_initializer,
)
out = dense_block(
"mlp_1",
out,
training,
channels=4096,
include_bn=False,
include_bias=True,
dropout_rate=0.5,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
beta_initializer=beta_initializer,
gamma_initializer=gamma_initializer,
)
logits = dense_block(
"mlp_2",
out,
training,
channels=num_classes,
include_bn=False,
include_bias=True,
act=final_act,
dropout_rate=0.5,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer,
beta_initializer=beta_initializer,
gamma_initializer=gamma_initializer,
)
else:
logits = x_tens
return logits
def dense_block(
name: str,
x_tens: tf_compat.Tensor,
training: Union[bool, tf_compat.Tensor],
channels: int,
include_bn: bool = False,
include_bias: bool = None,
dropout_rate: float = None,
act: Union[None, str] = "relu",
kernel_initializer=tf_compat.glorot_uniform_initializer(),
bias_initializer=tf_compat.zeros_initializer(),
beta_initializer=tf_compat.zeros_initializer(),
gamma_initializer=tf_compat.ones_initializer(),
):
"""
Create a dense or fully connected op and supporting ops
(batch norm, activation, etc) in the current graph and scope.
:param name: The name to group all ops under in the graph
:param x_tens: The input tensor to apply a fully connected and supporting ops to
:param training: A bool or tensor to indicate if the net is being run
in training mode or not. Used for batch norm and dropout
:param channels: The number of output channels from the dense op
:param include_bn: True to include a batch norm operation after the conv,
False otherwise
:param include_bias: If left unset, will add a bias if not include_bn.
Otherwise can be set to True to include a bias after the convolution,
False otherwise.
:param dropout_rate: The dropout rate to apply after the fully connected
and batch norm if included. If none, will not include batch norm
:param act: The activation to apply after the conv op and batch norm (if included).
Default is "relu", set to None for no activation.
:param kernel_initializer: The initializer to use for the fully connected kernels
:param bias_initializer: The initializer to use for the bias variable,
if a bias is included
:param beta_initializer: The initializer to use for the beta variable,
if batch norm is included
:param gamma_initializer: The initializer to use for the gamma variable,
if gamma is included
:return: the tensor after all ops have been applied
"""
if include_bias is None:
include_bias = not include_bn
with tf_compat.variable_scope(name, reuse=tf_compat.AUTO_REUSE):
out = tf_compat.layers.dense(
x_tens,
units=channels,
use_bias=include_bias,
kernel_initializer=kernel_initializer,
bias_initializer=bias_initializer if include_bias else None,
name="fc",
)
if include_bn:
out = tf_compat.layers.batch_normalization(
out,
axis=1,
momentum=BN_MOMENTUM,
epsilon=BN_EPSILON,
beta_initializer=beta_initializer,
gamma_initializer=gamma_initializer,
training=training,
name="bn",
)
if dropout_rate and dropout_rate > 0.0:
out = tf_compat.layers.dropout(out,
dropout_rate,
training=training,
name="dropout")
out = activation(out, act)
return out
