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