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 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 mobilenet(
inputs: tf_compat.Tensor,
training: Union[bool, tf_compat.Tensor] = True,
num_classes: int = 1000,
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:
"""
Standard MobileNet implementation with width=1.0;
expected input shape is (B, 224, 224, 3)
:param inputs: 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 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
"""
sec_settings = [
MobileNetSection(num_blocks=1, out_channels=64, downsample=False),
MobileNetSection(num_blocks=2, out_channels=128, downsample=True),
MobileNetSection(num_blocks=2, out_channels=256, downsample=True),
MobileNetSection(num_blocks=6, out_channels=512, downsample=True),
MobileNetSection(num_blocks=2, out_channels=1024, downsample=True),
]
return mobilenet_const(
inputs,
training,
sec_settings,
num_classes,
class_type,
kernel_initializer,
bias_initializer,
beta_initializer,
gamma_initializer,
)
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
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 mobilenet_v2_width(
width_mult: float,
inputs: tf_compat.Tensor,
training: Union[bool, tf_compat.Tensor] = True,
num_classes: int = 1000,
class_type: str = None,
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:
"""
Standard MobileNetV2 implementation for a given width;
expected input shape is (B, 224, 224, 3)
:param width_mult: The width multiplier for the architecture to create.
1.0 is standard, 0.5 is half the size, 2.0 is twice the size.
:param inputs: 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 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
"""
sec_settings = [
MobileNetV2Section(
num_blocks=1,
out_channels=16,
exp_channels=32,
downsample=False,
init_section=True,
width_mult=width_mult,
),
MobileNetV2Section(
num_blocks=2,
out_channels=24,
exp_ratio=6,
downsample=True,
init_section=False,
width_mult=width_mult,
),
MobileNetV2Section(
num_blocks=3,
out_channels=32,
exp_ratio=6,
downsample=True,
init_section=False,
width_mult=width_mult,
),
MobileNetV2Section(
num_blocks=4,
out_channels=64,
exp_ratio=6,
downsample=True,
init_section=False,
width_mult=width_mult,
),
MobileNetV2Section(
num_blocks=3,
out_channels=96,
exp_ratio=6,
downsample=False,
init_section=False,
width_mult=width_mult,
),
MobileNetV2Section(
num_blocks=3,
out_channels=160,
exp_ratio=6,
downsample=True,
init_section=False,
width_mult=width_mult,
),
MobileNetV2Section(
num_blocks=1,
out_channels=320,
exp_ratio=6,
downsample=False,
init_section=False,
width_mult=width_mult,
),
]
return mobilenet_v2_const(
inputs,
training,
sec_settings,
num_classes,
class_type,
kernel_initializer,
bias_initializer,
beta_initializer,
gamma_initializer,
)